module Ipopt.NLP where
import Ipopt.AnyRF
import Control.Applicative
import Control.Lens
import Control.Monad
import Control.Monad.Identity
import Control.Monad.State
import Data.Foldable (toList)
import Data.IntMap (IntMap)
import Data.List
import Data.Monoid
import Data.Monoid (First)
import Data.Sequence (Seq)
import Data.Vector (Vector)
import Foreign.C.Types (CDouble(..))
import qualified Data.Foldable as F
import qualified Data.IntMap as IM
import qualified Data.Map as M
import qualified Data.Sequence as Seq
import qualified Data.Set as S
import qualified Data.Vector as V
import qualified Data.Vector.Generic as VG
import Data.Vector ((!))
import qualified Data.Vector.Storable as VS
import Data.Maybe
import qualified Data.VectorSpace as VectorSpace
import Ipopt.Raw
data NLPFun = NLPFun
{ _funF, _funG :: AnyRF Seq,
_boundX, _boundG :: Seq (Double,Double) }
instance Show NLPFun where
show (NLPFun f g a b) = "NLPFun <f> <g> {" ++ show a ++ "}{" ++ show b ++ "}"
data NLPState = NLPState
{
_nMax :: Ix,
_currentEnv :: [String],
_variables :: M.Map String Ix,
_variablesInv :: IxMap String,
_constraintLabels, _objLabels :: IntMap String,
_varLabels :: IxMap String,
_varEnv :: IxMap (S.Set [String]),
_constraintEnv, _objEnv :: IntMap [String],
_nlpfun :: NLPFun,
_defaultBounds :: (Double,Double),
_initX :: Vector Double }
deriving (Show)
newtype IxMap a = IxMap (IntMap a)
deriving (Show, Monoid, Functor)
newtype Ix = Ix { _varIx :: Int } deriving Show
nlpstate0 = NLPState (Ix (1)) mempty mempty
mempty mempty mempty
mempty
mempty mempty mempty
(mempty :: NLPFun)
(1/0, 1/0)
V.empty
type NLPT = StateT NLPState
type NLP = NLPT IO
instance Monoid NLPFun where
NLPFun f g bx bg `mappend` NLPFun f' g' bx' bg' = NLPFun (f <> f') (g <> g') (bx <> bx') (bg <> bg')
mempty = NLPFun mempty mempty mempty mempty
makeLenses ''NLPFun
makeLenses ''Ix
makeLenses ''NLPState
makeIso ''IxMap
ix_ (Ix i) = from ixMap . ix i
at_ (Ix i) = from ixMap . at i
copyEnv to n = do
ev <- use currentEnv
cloneLens to %= IM.insert n ev
addDesc to Nothing n = return ()
addDesc to (Just x) n = do
to %= IM.insert n x
solveNLP' :: (VG.Vector v Double, MonadIO m) =>
(IpProblem -> IO ())
-> NLPT m (IpOptSolved v)
solveNLP' setOpts = do
(xl,xu) <- join $ uses (nlpfun . boundX) seqToVecs
(gl,gu) <- join $ uses (nlpfun . boundG) seqToVecs
AnyRF fs <- use (nlpfun . funF)
AnyRF gs <- use (nlpfun . funG)
p <- liftIO (createIpoptProblemAD xl xu gl gu (F.sum . fs) (V.fromList . toList . gs))
liftIO (setOpts p)
x0 <- uses initX (V.convert . V.map CDouble)
r <- liftIO (ipoptSolve p =<< VS.thaw x0)
liftIO $ freeIpoptProblem p
return r
addG :: Monad m
=> Maybe String
-> (Double,Double)
-> AnyRF Identity
-> NLPT m ()
addG d b (AnyRF f) = do
nlpfun . boundG %= (Seq.|> b)
nlpfun . funG %= \(AnyRF fs) -> AnyRF $ \x -> fs x Seq.|> runIdentity (f x)
n <- use (nlpfun . boundG . to Seq.length)
copyEnv constraintEnv n
addDesc constraintLabels d n
addF :: Monad m
=> Maybe String
-> AnyRF Identity
-> NLPT m ()
addF d (AnyRF f) = do
nlpfun . funF %= \(AnyRF fs) -> AnyRF $ \x -> fs x Seq.|> runIdentity (f x)
n <- use (objEnv . to ((+1) . IM.size))
copyEnv objEnv n
addDesc objLabels d n
var' :: (Monad m, Functor m)
=> Maybe (Double,Double)
-> Maybe String
-> String
-> NLPT m Ix
var' bs longDescription s = do
ev <- use currentEnv
m <- use variables
let s' = intercalate "." (reverse (s:ev))
n <- case M.lookup s' m of
Nothing -> do
nMax %= over varIx (+1)
db <- use defaultBounds
nlpfun . boundX %= (Seq.|> db)
n' <- use nMax
variables %= M.insert s' n'
variablesInv . at_ n' .= Just s'
F.for_ longDescription $ \d -> varLabels . at_ n' %= (<> Just d)
initX %= let x0 = fromMaybe 0 $
bs >>= \(a,b) -> find valid [(a+b)/2, a, b]
valid x = not $ isInfinite x || isNaN x
in (`V.snoc` x0)
return n'
Just n -> return n
varEnv . at_ n %= (<> Just (S.singleton ev))
F.for_ longDescription $ \str -> varLabels . at_ n %= (<> Just str)
F.traverse_ (narrowBounds n) bs
return n
var bs s = ixToVar <$> var' bs Nothing s
varFresh' :: (Monad m, Functor m) => Maybe (Double,Double) -> String -> NLPT m Ix
varFresh' bs s = do
existing <- gets (^? variables . ix s)
case existing of
Just a -> return a
Nothing -> do
m <- use variables
let n = M.size m + 1
Just sUniq <- return $ find (`M.notMember` m) $ s : iterate (++"_") (s ++ show n)
var' bs Nothing sUniq
varFresh bs s = fmap ixToVar $ varFresh' bs s
inEnv :: Monad m => String -> NLPT m a -> NLPT m a
inEnv s action = do
pushEnv s
r <- action
popEnv
return r
pushEnv :: Monad m => String -> NLPT m ()
pushEnv s = currentEnv %= (s:)
popEnv :: Monad m => NLPT m String
popEnv = do
n : ns <- use currentEnv
currentEnv .= ns
return n
splitVar :: (Monad m, Functor m)
=> Double
-> Ix
-> NLPT m (AnyRF Identity, AnyRF Identity)
splitVar b i = do
Just s <- gets (^? variablesInv . ix_ i)
let x = ixToVar i
xPlus <- varFresh' (Just (0, 1/0)) (s ++ "+")
xMinus <- varFresh' (Just (0, 1/0)) (s ++ "-")
x0 <- uses initX (! view varIx i)
initX %= (V.// [(view varIx xPlus, max 0 (x0 b)), (view varIx xMinus, max 0 ( b x0) )])
addG (Just ("splitVar: " ++ s)) (b, b) (x + ixToVar xMinus ixToVar xPlus)
return (ixToVar xMinus, ixToVar xPlus)
ixToVar :: Ix -> AnyRF Identity
ixToVar (Ix i) = AnyRF (\v -> Identity (v V.! i))
setBounds :: Monad m => Ix -> (Double,Double) -> NLPT m ()
setBounds (Ix i) bs = nlpfun . boundX %= Seq.update i bs
narrowBounds :: Monad m => Ix -> (Double,Double) -> NLPT m ()
narrowBounds (Ix i) (a,b) = nlpfun . boundX . ix i %= \(a',b') -> (max a a', min b b')
seqToVecs :: MonadIO m => Seq (Double,Double) -> m (Vec,Vec)
seqToVecs x = let (a,b) = unzip (toList x) in liftIO $ do
a' <- VS.thaw (VS.map CDouble (VS.fromList a))
b' <- VS.thaw (VS.map CDouble (VS.fromList b))
return (a',b')