diff --git a/README.md b/README.md --- a/README.md +++ b/README.md @@ -5,8 +5,9 @@ This library has a few distinct features, which may later be broken into separate packages: * high-level, strongly-typed interface to CasADi * NLP modeling/solving (examples/EasyNlp.hs, examples/BasicNlp.hs, examples/SofaExpando.hs) -* OCP modeling/solving (examles/Glider.hs, examples/DaeColl.hs) -* proof of concept monadic NLP/OCP modeling DSL (examples/{NlpDsl.hs,OcpDslRocket.hs}) +* OCP modeling/solving (examles/Glider.hs, examples/DaeColl.hs, examples/Rocket.hs, etc) +* toy OCP solver interface (examples/ToyOcp.hs) +* proof of concept monadic NLP modeling DSL (examples/NlpDsl.hs) * live plotter for OCP solving (examples/Dynoplot.hs) This package is built on top of CasADi (www.casadi.org). @@ -16,8 +17,8 @@ Please keep in mind that this library is continually evolving as my PhD progresses and I expect it to be very unstable. The API is also very messy as the library is evolving fast and it's unclear which parts are internal and external. -Specifically, the matrix and vector views (J and M) aren't polished enough and don't yet work without leaking internals. Nevertheless, I have started making hackage releases so that my few users have some snapshots to version-constrain against. +The library is tested on travis-ci, so the unit tests pass and the examples build. To install: diff --git a/dynobud.cabal b/dynobud.cabal --- a/dynobud.cabal +++ b/dynobud.cabal @@ -1,5 +1,5 @@ name: dynobud -version: 1.4.0.0 +version: 1.7.1.0 synopsis: your dynamic optimization buddy description: See readme at license: LGPL-3 @@ -19,21 +19,24 @@ library exposed-modules: Dyno.AutoScaling + Dyno.FormatTime Dyno.LagrangePolynomials Dyno.TypeVecs Dyno.MultipleShooting + Dyno.SimpleOcp Dyno.Ocp Dyno.OcpHomotopy + Dyno.DirectCollocation.ActiveConstraints Dyno.DirectCollocation.Dynamic Dyno.DirectCollocation.Export Dyno.DirectCollocation.Formulate + Dyno.DirectCollocation.FormulateCov Dyno.DirectCollocation.Integrate Dyno.DirectCollocation.Interpolate - Dyno.DirectCollocation.Profile Dyno.DirectCollocation.Quadratures Dyno.DirectCollocation.Robust Dyno.DirectCollocation.Types - Dyno.SXElement + Dyno.Integrate Dyno.View.Cov -- Dyno.View.CustomFunction Dyno.View.Fun @@ -54,6 +57,7 @@ Dyno.NlpSolver Dyno.NlpUtils Dyno.Solvers + Dyno.SolverInternal -- Dyno.Sqp.Sqp -- Dyno.Sqp.LineSearch @@ -69,15 +73,19 @@ vector-binary-instances, mtl >=2.2.1, containers >=0.5, - hmatrix, + hmatrix >= 0.17.0.1, + hmatrix-gsl >= 0.17.0.0, linear >= 1.3.1.1, reflection >= 1.3.2, + aeson, binary, cereal, distributive, process, Plot-ho-matic >= 0.5.0.2, - generic-accessors >= 0.1.0.1 + generic-accessors >= 0.4.2.0, + time, + directory -- cplex hs-source-dirs: src default-language: Haskell2010 @@ -102,6 +110,19 @@ , casadi-bindings ghc-options: -O2 +executable toy-ocp + if flag(examples) + Buildable: True + else + Buildable: False + hs-source-dirs: examples + main-is: ToyOcp.hs + default-language: Haskell2010 + build-depends: dynobud + , base >=4.6 && < 5 + ghc-options: -O2 + + executable multiple_shooting if flag(examples) Buildable: True @@ -134,7 +155,7 @@ default-language: Haskell2010 build-depends: dynobud, base >=4.6 && < 5, - not-gloss >= 0.7.0.1, + not-gloss >= 0.7.4.0, stm, containers, -- binary, @@ -374,6 +395,24 @@ cmdargs ghc-options: -O2 -with-rtsopts=-T +executable quadrature-test + if flag(examples) + Buildable: True + else + Buildable: False + hs-source-dirs: examples + main-is: Quadrature.hs + default-language: Haskell2010 + build-depends: dynobud, + base >=4.6 && < 5, + generic-accessors >= 0.1.0.0, + bytestring, + zeromq4-haskell, +-- binary, + cereal, + vector + ghc-options: -threaded -O2 + --test-suite lp_tests -- type: exitcode-stdio-1.0 -- hs-source-dirs: tests @@ -414,7 +453,8 @@ type: exitcode-stdio-1.0 hs-source-dirs: tests main-is: NewUnitTests.hs - other-modules: IntegrationTests + other-modules: QuadratureTests + IntegrationTests VectorizeTests ViewTests Utils @@ -434,3 +474,12 @@ hmatrix-gsl, base >=4.6 && < 5 ghc-options: -O2 + +test-suite doctests + type: exitcode-stdio-1.0 + main-is: Doctests.hs + build-depends: base >= 4 && < 5, + doctest >= 0.8 + default-language: Haskell2010 + ghc-options: -threaded + hs-source-dirs: tests diff --git a/examples/BasicNlp.hs b/examples/BasicNlp.hs --- a/examples/BasicNlp.hs +++ b/examples/BasicNlp.hs @@ -5,8 +5,8 @@ -- todo: comment up the multiple shooting code as an example {-# OPTIONS_GHC -Wall #-} -{-# Language DeriveFunctor #-} -{-# Language DeriveGeneric #-} +{-# LANGUAGE DeriveFunctor #-} +{-# LANGUAGE DeriveGeneric #-} module Main where diff --git a/examples/DaePendulum.hs b/examples/DaePendulum.hs --- a/examples/DaePendulum.hs +++ b/examples/DaePendulum.hs @@ -1,10 +1,10 @@ {-# OPTIONS_GHC -Wall #-} -{-# Language TypeFamilies #-} -{-# Language FlexibleInstances #-} -{-# Language DeriveFunctor #-} -{-# Language DeriveGeneric #-} -{-# Language DataKinds #-} -{-# Language PolyKinds #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE DeriveFunctor #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE PolyKinds #-} module Main where @@ -16,6 +16,7 @@ import Dyno.Vectorize import Dyno.View.View ( View(..), J ) +import Dyno.View.JV ( catJV ) import Dyno.Solvers import Dyno.Nlp import Dyno.NlpUtils @@ -37,6 +38,9 @@ type instance C PendOcp = PendBc type instance H PendOcp = None type instance Q PendOcp = None +type instance QO PendOcp = None +type instance FP PendOcp = None +type instance PO PendOcp = None data PendX a = PendX { pX :: a , pY :: a @@ -65,20 +69,20 @@ instance Lookup (PendO ()) instance Lookup (PendP ()) -mayer :: a -> PendX a -> PendX a -> None a -> PendP a -> a -mayer tf _ _ _ _ = tf +mayer :: a -> PendX a -> PendX a -> None a -> PendP a -> None a -> a +mayer tf _ _ _ _ _ = tf -lagrange :: Floating a => PendX a -> PendZ a -> PendU a -> PendP a -> PendO a -> a -> a -> a -lagrange _ _ u _ _ _ tf = 1e-3*torque'**2 / tf +lagrange :: Floating a => PendX a -> PendZ a -> PendU a -> PendP a -> None a -> PendO a -> a -> a -> a +lagrange _ _ u _ _ _ _ tf = 1e-3*torque'**2 / tf where PendU torque' = u r :: Floating a => a r = 0.3 -pendDae :: Floating a => PendX a -> PendX a -> PendZ a -> PendU a -> PendP a -> a -> (PendR a, PendO a) +pendDae :: Floating a => PendX a -> PendX a -> PendZ a -> PendU a -> PendP a -> None a -> a -> (PendR a, PendO a) pendDae (PendX x' y' vx' vy' torque') (PendX x y vx vy torque) - (PendZ tau) (PendU uTorque') (PendP m) _ = (residual, outputs) + (PendZ tau) (PendU uTorque') (PendP m) _ _ = (residual, outputs) where residual = PendR (x' - vx) (y' - vy) @@ -92,34 +96,45 @@ fy = -torque*x + m*9.8 pendOcp :: OcpPhase' PendOcp -pendOcp = OcpPhase { ocpMayer = mayer - , ocpLagrange = lagrange - , ocpQuadratures = \_ _ _ _ _ _ _ -> None - , ocpDae = pendDae - , ocpBc = bc - , ocpPathC = pathc - , ocpPathCBnds = None - , ocpBcBnds = bcBnds - , ocpXbnd = xbnd - , ocpUbnd = ubnd - , ocpZbnd = PendZ (Just (-200), Just 200) - , ocpPbnd = PendP (Just 0.3, Just 0.3) - , ocpTbnd = (Just 0.1, Just 5) - , ocpObjScale = Nothing - , ocpTScale = Nothing - , ocpXScale = Just pendXScale - , ocpZScale = Just (PendZ 10) - , ocpUScale = Just (PendU 50) - , ocpPScale = Just (PendP 0.3) - , ocpResidualScale = Nothing - , ocpBcScale = Just $ PendBc pendXScale pendXScale - , ocpPathCScale = Just None - } +pendOcp = + OcpPhase + { ocpMayer = mayer + , ocpLagrange = lagrange + , ocpQuadratures = \_ _ _ _ _ _ _ _ -> None + , ocpQuadratureOutputs = \_ _ _ _ _ _ _ _ -> None + , ocpDae = pendDae + , ocpBc = bc + , ocpPathC = pathc + , ocpPlotOutputs = \_ _ _ _ _ _ _ _ _ _ _ -> None + , ocpObjScale = Nothing + , ocpTScale = Nothing + , ocpXScale = Just pendXScale + , ocpZScale = Just (PendZ 10) + , ocpUScale = Just (PendU 50) + , ocpPScale = Just (PendP 0.3) + , ocpResidualScale = Nothing + , ocpBcScale = Just $ PendBc pendXScale pendXScale + , ocpPathCScale = Just None + } + +pendOcpInputs :: OcpPhaseInputs' PendOcp +pendOcpInputs = + OcpPhaseInputs + { ocpPathCBnds = None + , ocpBcBnds = bcBnds + , ocpXbnd = xbnd + , ocpUbnd = ubnd + , ocpZbnd = PendZ (Just (-200), Just 200) + , ocpPbnd = PendP (Just 0.3, Just 0.3) + , ocpTbnd = (Just 0.1, Just 5) + , ocpFixedP = None + } + pendXScale :: PendX Double pendXScale = PendX 0.3 0.3 1 1 10 -pathc :: Floating a => PendX a -> PendZ a -> PendU a -> PendP a -> PendO a -> a -> None a -pathc _ _ _ _ _ _ = None +pathc :: Floating a => PendX a -> PendZ a -> PendU a -> PendP a -> None a -> PendO a -> a -> None a +pathc _ _ _ _ _ _ _ = None xbnd :: PendX Bounds xbnd = PendX { pX = (Nothing, Nothing) @@ -132,8 +147,8 @@ ubnd :: PendU Bounds ubnd = PendU (Just (-100), Just 100) -bc :: Floating a => PendX a -> PendX a -> None a -> PendP a -> a -> PendBc a -bc x0 xf _ _ _ = PendBc x0 xf +bc :: Floating a => PendX a -> PendX a -> None a -> PendP a -> None a -> a -> PendBc a +bc x0 xf _ _ _ _ = PendBc x0 xf bcBnds :: PendBc Bounds bcBnds = @@ -185,12 +200,12 @@ main :: IO () main = do - cp <- makeCollProblem Legendre pendOcp guess + cp <- makeCollProblem Legendre pendOcp pendOcpInputs guess withCallback $ \send -> do let nlp = cpNlp cp meta = toMeta (cpMetaProxy cp) - cb' traj = do - plotPoints <- cpPlotPoints cp traj + cb' traj _ = do + plotPoints <- cpPlotPoints cp traj (catJV None) send (plotPoints, meta) _ <- solveNlp solver nlp (Just cb') -- _ <- solveNlp solver2 nlp Nothing diff --git a/examples/Dynoplot.hs b/examples/Dynoplot.hs --- a/examples/Dynoplot.hs +++ b/examples/Dynoplot.hs @@ -1,5 +1,5 @@ {-# OPTIONS_GHC -Wall #-} -{-# Language DeriveDataTypeable #-} +{-# LANGUAGE DeriveDataTypeable #-} module Main ( main ) where diff --git a/examples/EasyNlp.hs b/examples/EasyNlp.hs --- a/examples/EasyNlp.hs +++ b/examples/EasyNlp.hs @@ -2,8 +2,8 @@ -- the most basic NLP interface. {-# OPTIONS_GHC -Wall #-} -{-# Language DeriveFunctor #-} -{-# Language DeriveGeneric #-} +{-# LANGUAGE DeriveFunctor #-} +{-# LANGUAGE DeriveGeneric #-} module Main where diff --git a/examples/ExampleDsl/LogsAndErrors.hs b/examples/ExampleDsl/LogsAndErrors.hs --- a/examples/ExampleDsl/LogsAndErrors.hs +++ b/examples/ExampleDsl/LogsAndErrors.hs @@ -1,6 +1,6 @@ {-# OPTIONS_GHC -Wall #-} -{-# Language PackageImports #-} -{-# Language FlexibleContexts #-} +{-# LANGUAGE PackageImports #-} +{-# LANGUAGE FlexibleContexts #-} module ExampleDsl.LogsAndErrors ( ErrorMessage (..) diff --git a/examples/ExampleDsl/NlpMonad.hs b/examples/ExampleDsl/NlpMonad.hs --- a/examples/ExampleDsl/NlpMonad.hs +++ b/examples/ExampleDsl/NlpMonad.hs @@ -1,8 +1,8 @@ {-# OPTIONS_GHC -Wall #-} -{-# Language ScopedTypeVariables #-} -{-# Language PackageImports #-} -{-# Language GeneralizedNewtypeDeriving #-} -{-# Language RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE PackageImports #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE RankNTypes #-} module ExampleDsl.NlpMonad ( NlpMonad @@ -15,7 +15,7 @@ , solveStaticNlp ) where -import Control.Applicative ( Applicative ) +import qualified Control.Applicative as A import Control.Monad ( when ) import "mtl" Control.Monad.Reader ( MonadIO(..) ) import "mtl" Control.Monad.Except ( ExceptT, MonadError, runExceptT ) @@ -33,18 +33,20 @@ import Casadi.SharedObject ( soInit ) import Casadi.MX ( MX ) +import Casadi.SX ( SX ) import Casadi.SXFunction import Casadi.Function import Casadi.CMatrix ( veccat ) +import qualified Casadi.CMatrix as CM -import Dyno.View.Unsafe.View ( mkJ, unJ ) +import Dyno.View.Unsafe.View ( J(..), mkJ, unJ ) -import Dyno.SXElement ( SXElement, sxElementSym, sxElementToSX ) -import Dyno.Vectorize ( Id, fill ) +import Dyno.Vectorize ( Id, devectorize, fill ) import Dyno.TypeVecs ( Vec ) -import Dyno.View.View ( View(..), J, JNone(..), jfill ) +import Dyno.View.View ( View(..), JNone(..), jfill ) import Dyno.View.JV ( JV ) import Dyno.View.JVec ( JVec ) +import qualified Dyno.View.Symbolic as Sym import qualified Dyno.TypeVecs as TV import Dyno.Solvers ( Solver ) import Dyno.NlpUtils ( solveNlp ) @@ -53,6 +55,18 @@ import ExampleDsl.LogsAndErrors import ExampleDsl.Types +type SXElement = J (JV Id) SX + +sxElementSym :: String -> IO SXElement +sxElementSym = Sym.sym + +sxElementToSX :: SXElement -> SX +sxElementToSX (UnsafeJ x) + | (1,1) == sizes' = x + | otherwise = error $ "sxElementToSX: got non-scalar of size " ++ show sizes' + where + sizes' = (CM.size1 x, CM.size2 x) + --withEllipse :: Int -> String -> String --withEllipse n blah -- | length blah <= n = blah @@ -62,7 +76,7 @@ NlpMonad { runNlp :: ExceptT ErrorMessage (WriterT [LogMessage] (StateT NlpMonadState IO)) a } deriving ( Functor - , Applicative + , A.Applicative , Monad , MonadError ErrorMessage , MonadState NlpMonadState @@ -147,7 +161,7 @@ toG :: Dim ng => S.Seq (Constraint SXElement) -> Vec ng (SXElement, Bounds) -toG nlpConstraints' = TV.mkVec $ V.fromList $ F.toList $ fmap constr nlpConstraints' +toG nlpConstraints' = devectorize $ V.fromList $ F.toList $ fmap constr nlpConstraints' buildNlp :: forall nx ng . (Dim nx, Dim ng) => NlpMonadState -> IO (Nlp (JVec nx (JV Id)) JNone (JVec ng (JV Id)) MX) @@ -192,9 +206,9 @@ forall r . NlpMonad () -> Maybe (Vector Double -> IO Bool) -> M.Map String Double -> (forall x g . (View x, View g) - => Nlp x JNone g MX -> Maybe (J x (Vector Double) -> IO Bool) -> NlpMonadState -> IO r) + => Nlp x JNone g MX -> Maybe (J x (Vector Double) -> J JNone (Vector Double) -> IO Bool) -> NlpMonadState -> IO r) -> IO r -reifyNlp nlpmonad cb x0map f = do +reifyNlp nlpmonad cb0 x0map f = do (ret,logs,state) <- build nlpmonad case ret of Right _ -> return () @@ -211,9 +225,13 @@ TV.reifyDim ng $ \(Proxy :: Proxy ng) -> do nlp0 <- buildNlp state :: IO (Nlp (JVec nx (JV Id)) JNone (JVec ng (JV Id)) MX) let nlp = nlp0 { nlpX0 = mkJ x0 } - f nlp (fmap (. unJ) cb) state + cb = case cb0 of + Nothing -> Nothing + Just cb' -> Just $ \x _ -> cb' (unJ x) + f nlp cb state + solveStaticNlp :: Solver -> NlpMonad () -> [(String,Double)] -> Maybe (Vector Double -> IO Bool) @@ -226,8 +244,8 @@ show (xx,yy) foo :: - (View x, View p, View g) => - Nlp x p g MX -> Maybe (J x (Vector Double) -> IO Bool) -> NlpMonadState -> + (View x, View g) => + Nlp x JNone g MX -> Maybe (J x (Vector Double) -> J JNone (Vector Double) -> IO Bool) -> NlpMonadState -> IO (Either String String, Double, [(String,Double)]) foo nlp' cb' state = do (ret,nlpOut) <- solveNlp solverStuff nlp' cb' diff --git a/examples/ExampleDsl/Types.hs b/examples/ExampleDsl/Types.hs --- a/examples/ExampleDsl/Types.hs +++ b/examples/ExampleDsl/Types.hs @@ -19,9 +19,11 @@ import qualified Data.Sequence as S import qualified Data.Map as M import Control.Lens -import Data.Functor ( (<$>) ) -import Dyno.SXElement ( SXElement ) +import Casadi.SX ( SX ) +import Dyno.View.View ( J ) +import Dyno.View.JV ( JV ) +import Dyno.Vectorize ( Id ) data Constraint a = Eq2 a a | Ineq2 a a @@ -30,6 +32,8 @@ data Objective a = ObjectiveUnset | Objective a data HomotopyParam a = HomotopyParamUnset | HomotopyParam a +type SXElement = J (JV Id) SX + data NlpMonadState = NlpMonadState { nlpX :: S.Seq (String, SXElement) @@ -57,30 +61,30 @@ --makeLenses ''DaeState daeXDot :: Lens' DaeState (S.Seq (String, SXElement)) daeXDot f (DaeState xdot' x z u p o ss c) = - (\xdot -> DaeState xdot x z u p o ss c) <$> f xdot' + (\xdot -> DaeState xdot x z u p o ss c) `fmap` f xdot' {-# INLINE daeXDot #-} daeX :: Lens' DaeState (S.Seq (String, SXElement)) daeX f (DaeState xdot x' z u p o ss c) = - (\x -> DaeState xdot x z u p o ss c) <$> f x' + (\x -> DaeState xdot x z u p o ss c) `fmap` f x' {-# INLINE daeX #-} daeZ :: Lens' DaeState (S.Seq (String, SXElement)) daeZ f (DaeState xdot x z' u p o ss c) = - (\z -> DaeState xdot x z u p o ss c) <$> f z' + (\z -> DaeState xdot x z u p o ss c) `fmap` f z' {-# INLINE daeZ #-} daeU :: Lens' DaeState (S.Seq (String, SXElement)) daeU f (DaeState xdot x z u' p o ss c) = - (\u -> DaeState xdot x z u p o ss c) <$> f u' + (\u -> DaeState xdot x z u p o ss c) `fmap` f u' {-# INLINE daeU #-} daeP :: Lens' DaeState (S.Seq (String, SXElement)) daeP f (DaeState xdot x z u p' o ss c) = - (\p -> DaeState xdot x z u p o ss c) <$> f p' + (\p -> DaeState xdot x z u p o ss c) `fmap` f p' {-# INLINE daeP #-} daeO :: Lens' DaeState (M.Map String SXElement) daeO f (DaeState xdot x z u p o' ss c) = - (\o -> DaeState xdot x z u p o ss c) <$> f o' + (\o -> DaeState xdot x z u p o ss c) `fmap` f o' {-# INLINE daeO #-} diff --git a/examples/Glider.hs b/examples/Glider.hs --- a/examples/Glider.hs +++ b/examples/Glider.hs @@ -1,6 +1,6 @@ {-# OPTIONS_GHC -Wall #-} -{-# Language TypeFamilies #-} -{-# Language DataKinds #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE DataKinds #-} module Main ( main ) where @@ -9,6 +9,7 @@ import Dyno.Vectorize import Dyno.View.View +import Dyno.View.JV ( catJV ) import Dyno.Solvers --import Dyno.Sqp.Sqp --import Dyno.Sqp.LineSearch @@ -39,12 +40,15 @@ type instance C GliderOcp = AcX type instance H GliderOcp = None type instance Q GliderOcp = None +type instance QO GliderOcp = None +type instance FP GliderOcp = None +type instance PO GliderOcp = None -mayer :: Floating a => a -> AcX a -> AcX a -> None a -> None a -> a -mayer _ _ _ _ _ = 0 +mayer :: Floating a => a -> AcX a -> AcX a -> None a -> None a -> None a -> a +mayer _ _ _ _ _ _ = 0 -lagrange :: Floating a => AcX a -> None a -> AcU a -> None a -> None a -> a -> a -> a -lagrange (AcX _ _ _ _ (AcU surfs)) _ (AcU surfs') _ _ _ _ = +lagrange :: Floating a => AcX a -> None a -> AcU a -> None a -> None a -> None a -> a -> a -> a +lagrange (AcX _ _ _ _ (AcU surfs)) _ (AcU surfs') _ _ _ _ _ = elev**2 + rudd**2 + ail**2 + flaps**2 + 100*(elev'**2 + rudd'**2 + ail'**2 + flaps'**2) where @@ -58,8 +62,8 @@ ail' = csElev surfs' flaps' = csFlaps surfs' -dae :: Floating a => AcX a -> AcX a -> None a -> AcU a -> None a -> a -> (AcX a, None a) -dae x' x _ u _ _ = (aircraftDae (mass, inertia) fcs mcs refs x' x u, None) +dae :: Floating a => AcX a -> AcX a -> None a -> AcU a -> None a -> None a -> a -> (AcX a, None a) +dae x' x _ u _ _ _ = (aircraftDae (mass, inertia) fcs mcs refs x' x u, None) where mass = bettyMass inertia = bettyInertia @@ -70,18 +74,12 @@ ocp :: OcpPhase' GliderOcp ocp = OcpPhase { ocpMayer = mayer , ocpLagrange = lagrange - , ocpQuadratures = \_ _ _ _ _ _ _ -> None + , ocpQuadratures = \_ _ _ _ _ _ _ _ -> None + , ocpQuadratureOutputs = \_ _ _ _ _ _ _ _ -> None , ocpDae = dae , ocpBc = bc , ocpPathC = pathc - , ocpPathCBnds = None - , ocpBcBnds = fill (Just 0, Just 0) - , ocpXbnd = xbnd - , ocpUbnd = ubnd - , ocpZbnd = None - , ocpPbnd = None - , ocpTbnd = (Just 0.5, Just 0.5) --- , ocpTbnd = (Just 4, Just 4) + , ocpPlotOutputs = \_ _ _ _ _ _ _ _ _ _ _ -> None , ocpObjScale = Nothing , ocpTScale = Nothing , ocpXScale = Nothing @@ -93,9 +91,23 @@ , ocpPathCScale = Nothing } -pathc :: x a -> z a -> u a -> p a -> None a -> a -> None a -pathc _ _ _ _ _ _ = None +ocpInputs :: OcpPhaseInputs' GliderOcp +ocpInputs = + OcpPhaseInputs + { ocpPathCBnds = None + , ocpBcBnds = fill (Just 0, Just 0) + , ocpXbnd = xbnd + , ocpUbnd = ubnd + , ocpZbnd = None + , ocpPbnd = None + , ocpTbnd = (Just 0.5, Just 0.5) +-- , ocpTbnd = (Just 4, Just 4) + , ocpFixedP = None + } +pathc :: x a -> z a -> u a -> p a -> None a -> None a -> a -> None a +pathc _ _ _ _ _ _ _ = None + xbnd :: AcX (Maybe Double, Maybe Double) xbnd = AcX { ac_r_n2b_n = fill (Nothing, Nothing) , ac_v_bn_b = fill (Nothing,Nothing) @@ -116,8 +128,8 @@ , csFlaps = (Just (d2r (-0.01)), Just (d2r 0.01)) } -bc :: Floating a => AcX a -> AcX a -> None a -> None a -> a -> AcX a -bc (AcX x0 v0 dcm0 w0 cs) _ _ _ _ = AcX x0 (v0 - V3 30 0 0) (dcm0 - eye3') w0 cs +bc :: Floating a => AcX a -> AcX a -> None a -> None a -> None a -> a -> AcX a +bc (AcX x0 v0 dcm0 w0 cs) _ _ _ _ _ = AcX x0 (v0 - V3 30 0 0) (dcm0 - eye3') w0 cs eye3' :: Num a => M33 a eye3' = @@ -130,13 +142,13 @@ main :: IO () main = do let guess = jfill 1 :: J (CollTraj' GliderOcp NCollStages CollDeg) (Vector Double) - cp <- makeCollProblem Legendre ocp guess + cp <- makeCollProblem Legendre ocp ocpInputs guess let nlp = cpNlp cp withCallback $ \send -> do let meta = toMeta (cpMetaProxy cp) - cb' traj = do - plotPoints <- cpPlotPoints cp traj + cb' traj _ = do + plotPoints <- cpPlotPoints cp traj (catJV None) send (plotPoints, meta) (msg,_) <- solveNlp ipoptSolver nlp (Just cb') diff --git a/examples/Glider/AeroCoeffs.hs b/examples/Glider/AeroCoeffs.hs --- a/examples/Glider/AeroCoeffs.hs +++ b/examples/Glider/AeroCoeffs.hs @@ -1,13 +1,13 @@ {-# OPTIONS_GHC -Wall #-} -{-# Language DeriveFunctor #-} -{-# Language DeriveFoldable #-} -{-# Language DeriveGeneric #-} +{-# LANGUAGE DeriveFunctor #-} +{-# LANGUAGE DeriveFoldable #-} +{-# LANGUAGE DeriveGeneric #-} module Glider.AeroCoeffs where import GHC.Generics ( Generic, Generic1 ) -import Data.Foldable ( Foldable ) +import qualified Data.Foldable as F import Linear import Accessors ( Lookup ) @@ -22,7 +22,7 @@ , csRudder :: a , csAil :: a , csFlaps :: a - } deriving (Eq, Functor, Foldable, Generic, Generic1, Show) + } deriving (Eq, Functor, F.Foldable, Generic, Generic1, Show) instance Vectorize ControlSurfaces instance (Lookup a, Generic a) => Lookup (ControlSurfaces a) diff --git a/examples/Glider/Aircraft.hs b/examples/Glider/Aircraft.hs --- a/examples/Glider/Aircraft.hs +++ b/examples/Glider/Aircraft.hs @@ -1,7 +1,7 @@ {-# OPTIONS_GHC -Wall #-} -{-# Language ScopedTypeVariables #-} -{-# Language DeriveFunctor #-} -{-# Language DeriveGeneric #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE DeriveFunctor #-} +{-# LANGUAGE DeriveGeneric #-} module Glider.Aircraft ( AcX(..), AcU(..), aircraftDae ) where diff --git a/examples/Homotopy.hs b/examples/Homotopy.hs --- a/examples/Homotopy.hs +++ b/examples/Homotopy.hs @@ -1,6 +1,6 @@ {-# OPTIONS_GHC -Wall #-} -{-# Language DeriveFunctor #-} -{-# Language DeriveGeneric #-} +{-# LANGUAGE DeriveFunctor #-} +{-# LANGUAGE DeriveGeneric #-} module Main where @@ -87,5 +87,5 @@ printf "X: (%.3f,%.3f), P: (%.3f, %.3f), a: %.4f\n" x y px py alpha return () pfs = [catJV (P 2 0), catJV (P 3 0)] - opt <- solveNlpHomotopy 1e-3 hp solver Nothing myNlp pfs Nothing (Just cbp) + opt <- solveNlpHomotopy 1e-3 hp solver myNlp pfs Nothing (Just cbp) print opt diff --git a/examples/MultipleShooting.hs b/examples/MultipleShooting.hs --- a/examples/MultipleShooting.hs +++ b/examples/MultipleShooting.hs @@ -1,9 +1,9 @@ {-# OPTIONS_GHC -Wall #-} -{-# Language ScopedTypeVariables #-} -{-# Language DeriveGeneric #-} -{-# Language DeriveFunctor #-} -{-# Language DataKinds #-} -{-# Language PolyKinds #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE DeriveFunctor #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE PolyKinds #-} module Main ( main @@ -11,10 +11,7 @@ import GHC.Generics ( Generic, Generic1 ) -import qualified Data.Vector as V import qualified Data.Foldable as F -import Control.Applicative ( Applicative(..) ) -import Linear import Graphics.Rendering.Chart hiding ( x0 ) import Graphics.Rendering.Chart.Gtk @@ -43,16 +40,6 @@ instance Vectorize X instance Vectorize U instance Vectorize P -instance Applicative X where - pure = fill - x0 <*> x1 = devectorize (V.zipWith id (vectorize x0) (vectorize x1)) -instance Applicative U where - pure = fill - x0 <*> x1 = devectorize (V.zipWith id (vectorize x0) (vectorize x1)) -instance Additive X where - zero = fill 0 -instance Additive U where - zero = fill 0 -- ocp specification ocp :: MsOcp X U P diff --git a/examples/NlpSolverEx.hs b/examples/NlpSolverEx.hs --- a/examples/NlpSolverEx.hs +++ b/examples/NlpSolverEx.hs @@ -1,8 +1,8 @@ -- | Example of NlpSolver monad and autoscaling {-# OPTIONS_GHC -Wall #-} -{-# Language DeriveFunctor #-} -{-# Language DeriveGeneric #-} +{-# LANGUAGE DeriveFunctor #-} +{-# LANGUAGE DeriveGeneric #-} module Main where diff --git a/examples/Quadrature.hs b/examples/Quadrature.hs new file mode 100644 --- /dev/null +++ b/examples/Quadrature.hs @@ -0,0 +1,217 @@ +{-# OPTIONS_GHC -Wall #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE DeriveFunctor #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE FlexibleInstances #-} + +module Main + ( main + , lolMahQ + ) where + +import GHC.Generics ( Generic, Generic1 ) + +import Data.Vector ( Vector ) +import Text.Printf ( printf ) + +import Accessors ( Lookup ) + +import Dyno.Vectorize ( Vectorize(..), None(..), Id(..) ) +import Dyno.View.View ( View(..), J ) +import Dyno.View.JV ( splitJV, catJV ) +import Dyno.Solvers +import Dyno.Nlp ( NlpOut(..), Bounds ) +import Dyno.NlpUtils +import Dyno.Ocp +import Dyno.DirectCollocation.Formulate +import Dyno.DirectCollocation.Types +--import Dyno.DirectCollocation.Types ( CollTraj(..) ) +import Dyno.DirectCollocation.Dynamic ( toMeta ) +import Dyno.DirectCollocation.Quadratures ( QuadratureRoots(..) ) + +import Dynoplot.Callback ( withCallback ) + + +data QuadOcp +type instance X QuadOcp = QuadX +type instance Z QuadOcp = QuadZ +type instance U QuadOcp = QuadU +type instance P QuadOcp = QuadP +type instance R QuadOcp = QuadR +type instance O QuadOcp = QuadO +type instance C QuadOcp = QuadBc +type instance H QuadOcp = None +type instance Q QuadOcp = QuadQ +type instance QO QuadOcp = None +type instance FP QuadOcp = None +type instance PO QuadOcp = None + +data QuadX a = QuadX { xP :: a + , xV :: a + } deriving (Functor, Generic, Generic1, Show) +data QuadZ a = QuadZ deriving (Functor, Generic, Generic1, Show) +data QuadU a = QuadU deriving (Functor, Generic, Generic1, Show) +data QuadP a = QuadP deriving (Functor, Generic, Generic1, Show) +data QuadR a = QuadR (QuadX a) deriving (Functor, Generic, Generic1, Show) +data QuadO a = QuadO a deriving (Functor, Generic, Generic1, Show) +data QuadBc a = QuadBc (QuadX a) deriving (Functor, Generic, Generic1, Show) +data QuadQ a = QuadQ { lolMahQ :: a } deriving (Functor, Generic, Generic1, Show) + +instance Vectorize QuadX +instance Vectorize QuadZ +instance Vectorize QuadU +instance Vectorize QuadP +instance Vectorize QuadR +instance Vectorize QuadO +instance Vectorize QuadBc +instance Vectorize QuadQ + +instance Lookup (QuadX ()) +instance Lookup (QuadZ ()) +instance Lookup (QuadU ()) +instance Lookup (QuadO ()) +instance Lookup (QuadP ()) +instance Lookup (QuadQ ()) + +mayer :: Num a => QuadOrLagrange -> a -> QuadX a -> QuadX a -> QuadQ a -> QuadP a -> None a -> a +mayer TestQuadratures _ _ _ (QuadQ qf) _ _ = qf +mayer TestLagrangeTerm _ _ _ _ _ _ = 0 + +data QuadOrLagrange = TestQuadratures | TestLagrangeTerm deriving Show +data StateOrOutput = TestState | TestOutput deriving Show + +lagrange :: Num a => StateOrOutput -> QuadOrLagrange -> QuadX a -> QuadZ a -> QuadU a -> QuadP a -> None a -> QuadO a -> a -> a -> a +lagrange _ TestQuadratures _ _ _ _ _ _ _ _ = 0 +lagrange TestState TestLagrangeTerm (QuadX _ v) _ _ _ _ _ _ _ = v +lagrange TestOutput TestLagrangeTerm _ _ _ _ _ (QuadO v) _ _ = v + +quadratures :: Floating a => + StateOrOutput -> QuadX a -> QuadZ a -> QuadU a -> QuadP a -> None a -> QuadO a -> a -> a -> QuadQ a +quadratures TestState (QuadX _ v) _ _ _ _ _ _ _ = QuadQ v +quadratures TestOutput _ _ _ _ _ (QuadO v) _ _ = QuadQ v + +dae :: Floating a => QuadX a -> QuadX a -> QuadZ a -> QuadU a -> QuadP a -> None a -> a -> (QuadR a, QuadO a) +dae (QuadX p' v') (QuadX _ v) _ _ _ _ _ = (residual, outputs) + where + residual = + QuadR + QuadX { xP = p' - v + , xV = v' - alpha + } + outputs = QuadO v + +alpha :: Fractional a => a +alpha = 1 + +tf :: Fractional a => a +tf = 10.0 + +quadOcp :: StateOrOutput -> QuadOrLagrange -> OcpPhase' QuadOcp +quadOcp stateOrOutput quadOrLag = + OcpPhase + { ocpMayer = mayer quadOrLag + , ocpLagrange = lagrange stateOrOutput quadOrLag + , ocpQuadratures = quadratures stateOrOutput + , ocpQuadratureOutputs = \_ _ _ _ _ _ _ _ -> None + , ocpDae = dae + , ocpBc = bc + , ocpPathC = pathc + , ocpPlotOutputs = \_ _ _ _ _ _ _ _ _ _ _ -> None + , ocpObjScale = Nothing + , ocpTScale = Nothing + , ocpXScale = Nothing + , ocpZScale = Nothing + , ocpUScale = Nothing + , ocpPScale = Nothing + , ocpResidualScale = Nothing + , ocpBcScale = Nothing + , ocpPathCScale = Just None + } + +quadOcpInputs :: OcpPhaseInputs' QuadOcp +quadOcpInputs = + OcpPhaseInputs + { ocpPathCBnds = None + , ocpBcBnds = bcBnds + , ocpXbnd = xbnd + , ocpUbnd = ubnd + , ocpZbnd = QuadZ + , ocpPbnd = QuadP + , ocpTbnd = (Just tf, Just tf) + , ocpFixedP = None + } + +pathc :: Floating a => QuadX a -> QuadZ a -> QuadU a -> QuadP a -> None a -> QuadO a -> a -> None a +pathc _ _ _ _ _ _ _ = None + +xbnd :: QuadX Bounds +xbnd = QuadX { xP = (Nothing, Nothing) + , xV = (Nothing, Nothing) + } + +ubnd :: QuadU Bounds +ubnd = QuadU + +bc :: Floating a => QuadX a -> QuadX a -> QuadQ a -> QuadP a -> None a -> a -> QuadBc a +bc x0 _ _ _ _ _ = QuadBc x0 + +bcBnds :: QuadBc Bounds +bcBnds = + QuadBc + (QuadX + { xP = (Just 0, Just 0) + , xV = (Just 0, Just 0) + }) + +type NCollStages = 10 +type CollDeg = 3 + +guess :: QuadratureRoots -> J (CollTraj' QuadOcp NCollStages CollDeg) (Vector Double) +guess roots = cat $ makeGuess roots tf guessX guessZ guessU parm + where + guessX _ = QuadX { xP = 0 + , xV = 0 + } + guessZ _ = QuadZ + guessU _ = QuadU + parm = QuadP + + + +solver :: Solver +solver = ipoptSolver { options = [ ("expand", Opt True) + , ("linear_solver", Opt "ma86") + , ("ma86_order", Opt "metis") +-- , ("print_level", Opt (0 :: Int)) +-- , ("print_time", Opt False) + ]} + +goodSolution :: NlpOut + (CollTraj QuadX QuadZ QuadU QuadP NCollStages CollDeg) + (CollOcpConstraints QuadX QuadR QuadBc None NCollStages CollDeg) + (Vector Double) + -> String +goodSolution out = msg + where + msg = printf " objective: %.4f, expected: %.4f" f fExpected + fExpected = 0.5 * alpha * tf**2 :: Double + Id f = splitJV (fOpt out) + +compareIntegration :: (QuadratureRoots, StateOrOutput, QuadOrLagrange) -> IO () +compareIntegration (roots, stateOrOutput, quadOrLag) = do + withCallback $ \send -> do + cp <- makeCollProblem roots (quadOcp stateOrOutput quadOrLag) quadOcpInputs (guess roots) + let nlp = cpNlp cp + meta = toMeta (cpMetaProxy cp) + cb traj _ = do + plotPoints <- cpPlotPoints cp traj (catJV None) + send (plotPoints, meta) + (ret, out) <- solveNlp solver nlp (Just cb) + case ret of + Left msg -> return (error msg) + Right _ -> putStrLn (goodSolution out) + +main :: IO () +main = do + compareIntegration (Legendre, TestState, TestLagrangeTerm) diff --git a/examples/Rocket.hs b/examples/Rocket.hs --- a/examples/Rocket.hs +++ b/examples/Rocket.hs @@ -1,8 +1,8 @@ {-# OPTIONS_GHC -Wall #-} -{-# Language TypeFamilies #-} -{-# Language DeriveFunctor #-} -{-# Language DeriveGeneric #-} -{-# Language DataKinds #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE DeriveFunctor #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE DataKinds #-} module Main ( main ) where @@ -13,11 +13,13 @@ import Accessors ( Lookup ) import Dyno.View.View ( J, jfill ) -import Dyno.Nlp ( Bounds ) +import Dyno.View.JV ( catJV ) +import Dyno.Nlp ( NlpOut(..), Bounds ) import Dyno.Ocp import Dyno.Vectorize ( Vectorize, None(..), fill ) import Dyno.Solvers ( Solver(..), Opt(..), ipoptSolver ) import Dyno.NlpUtils ( solveNlp ) +import Dyno.DirectCollocation.ActiveConstraints import Dyno.DirectCollocation.Formulate ( CollProblem(..), makeCollProblem ) import Dyno.DirectCollocation.Types ( CollTraj' ) import Dyno.DirectCollocation.Dynamic ( toMeta ) @@ -29,11 +31,27 @@ OcpPhase { ocpMayer = mayer , ocpLagrange = lagrange - , ocpQuadratures = \_ _ _ _ _ _ _ -> None + , ocpQuadratures = \_ _ _ _ _ _ _ _ -> None + , ocpQuadratureOutputs = \_ _ _ _ _ _ _ _ -> None , ocpDae = dae , ocpBc = bc , ocpPathC = pathC - , ocpBcBnds = bcBnds + , ocpPlotOutputs = \_ _ _ _ _ _ _ _ _ _ _ -> None + , ocpObjScale = Nothing + , ocpTScale = Nothing + , ocpXScale = Nothing + , ocpZScale = Nothing + , ocpUScale = Nothing + , ocpPScale = Nothing + , ocpResidualScale = Nothing + , ocpBcScale = Nothing + , ocpPathCScale = Nothing + } + +rocketOcpInputs :: OcpPhaseInputs' RocketOcp +rocketOcpInputs = + OcpPhaseInputs + { ocpBcBnds = bcBnds , ocpPathCBnds = pathCBnds , ocpXbnd = RocketX { xPos = (Just 0, Nothing) @@ -46,15 +64,7 @@ { uThrustDot = (Just (-100), Just 100) } , ocpPbnd = fill (Nothing, Nothing) , ocpTbnd = (Just 4, Just 4) - , ocpObjScale = Nothing - , ocpTScale = Nothing - , ocpXScale = Nothing - , ocpZScale = Nothing - , ocpUScale = Nothing - , ocpPScale = Nothing - , ocpResidualScale = Nothing - , ocpBcScale = Nothing - , ocpPathCScale = Nothing + , ocpFixedP = None } data RocketOcp @@ -67,6 +77,9 @@ type instance P RocketOcp = None type instance Z RocketOcp = None type instance Q RocketOcp = None +type instance QO RocketOcp = None +type instance FP RocketOcp = None +type instance PO RocketOcp = None data RocketX a = RocketX @@ -102,9 +115,9 @@ dae :: Floating a - => RocketX a -> RocketX a -> None a -> RocketU a -> None a -> a + => RocketX a -> RocketX a -> None a -> RocketU a -> None a -> None a -> a -> (RocketX a, RocketO a) -dae (RocketX p' v' m' thrust') (RocketX _ v m thrust) _ (RocketU uThrust') _ _ = +dae (RocketX p' v' m' thrust') (RocketX _ v m thrust) _ (RocketU uThrust') _ _ _ = (residual, outputs) where residual = RocketX @@ -119,8 +132,8 @@ force = thrust - m*g -bc :: RocketX a -> RocketX a -> None a -> None a -> a -> RocketBc a -bc x0 xf _ _ _ = RocketBc x0 xf +bc :: RocketX a -> RocketX a -> None a -> None a -> None a -> a -> RocketBc a +bc x0 xf _ _ _ _ = RocketBc x0 xf bcBnds :: RocketBc Bounds bcBnds = @@ -129,18 +142,18 @@ , bcXF = RocketX (Just 0, Just 0) (Just 0, Just 0) (Nothing, Nothing) (Nothing, Nothing) } -mayer :: Floating a => a -> RocketX a -> RocketX a -> None a -> None a -> a -mayer _endTime _ (RocketX _ _ mf _) _ _ = -mf -- endTime +mayer :: Floating a => a -> RocketX a -> RocketX a -> None a -> None a -> None a -> a +mayer _endTime _ (RocketX _ _ mf _) _ _ _ = -mf -- endTime -pathC :: Floating a => RocketX a -> None a -> RocketU a -> None a -> RocketO a -> a -> RocketPathC a -pathC _ _ _ _ _ = RocketPathC +pathC :: Floating a => RocketX a -> None a -> RocketU a -> None a -> None a -> RocketO a -> a -> RocketPathC a +pathC _ _ _ _ _ _ = RocketPathC pathCBnds :: RocketPathC Bounds pathCBnds = RocketPathC (Nothing, Just 4) -lagrange :: Fractional a => RocketX a -> None a -> RocketU a -> None a -> RocketO a -> a -> a -> a -lagrange _ _ (RocketU u') _ _ _ _ = 1e-4*u'*u' +lagrange :: Fractional a => RocketX a -> None a -> RocketU a -> None a -> None a -> RocketO a -> a -> a -> a +lagrange _ _ (RocketU u') _ _ _ _ _ = 1e-4*u'*u' -- (1e-8*u*u + 1e-9*p*p + 1e-9*v*v + 1e-9*m*m) -- (1e-6*u*u + 1e-6*p*p + 1e-6*v*v + 1e-6*m*m) @@ -157,13 +170,20 @@ main = withCallback $ \send -> do - cp <- makeCollProblem Legendre rocketOcp guess + cp <- makeCollProblem Legendre rocketOcp rocketOcpInputs guess let nlp = cpNlp cp meta = toMeta (cpMetaProxy cp) - cb' traj = do - plotPoints <- cpPlotPoints cp traj + cb' traj _ = do + plotPoints <- cpPlotPoints cp traj (catJV None) send (plotPoints, meta) - _ <- solveNlp solver nlp (Just cb') - return () + (ret, opt) <- solveNlp solver nlp (Just cb') + case ret of + Left msg -> putStrLn $ "\nsolve failed with " ++ show msg + Right msg -> do + putStrLn $ "\nsolve succeeded with " ++ show msg + activeConstraints <- getActiveConstraints (cpConstraints cp) rocketOcp 1e-3 + (xOpt opt) (catJV None) rocketOcpInputs + putStrLn "\nactive constriants:" + putStr $ summarizeActiveConstraints activeConstraints diff --git a/examples/Sailboat.hs b/examples/Sailboat.hs --- a/examples/Sailboat.hs +++ b/examples/Sailboat.hs @@ -1,14 +1,14 @@ -- This example is based on Fabian Wierer's final project for -- Optimal Control and Estimation, 2014, taught by Prof. Moritz Diehl --- \Used with permission. +-- Used with permission. {-# OPTIONS_GHC -Wall #-} -{-# Language TypeFamilies #-} -{-# Language ScopedTypeVariables #-} -{-# Language FlexibleInstances #-} -{-# Language DeriveFunctor #-} -{-# Language DeriveGeneric #-} -{-# Language DataKinds #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE DeriveFunctor #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE DataKinds #-} module Main ( main , SbX(..) -- to suppress warnings about unused record names @@ -32,7 +32,7 @@ import Dyno.Vectorize import Dyno.View.View ( View(..), J ) -import Dyno.View.JV ( splitJV ) +import Dyno.View.JV ( catJV, splitJV ) import Dyno.Solvers import Dyno.NlpUtils import Dyno.Nlp ( NlpOut(..) ) @@ -52,6 +52,9 @@ type instance C SailboatOcp = SbBc type instance H SailboatOcp = None type instance Q SailboatOcp = None +type instance QO SailboatOcp = None +type instance FP SailboatOcp = None +type instance PO SailboatOcp = None data SbX a = SbX { xGamma :: a , xP :: V2 a @@ -101,7 +104,7 @@ clift :: Floating a => a -> a clift alpha = 2*pi*alpha*10/12 - exp (alpha/pi*180 - 12) + exp (-alpha/pi*180 - 12) -sbDae :: forall a . Floating a => SbX a -> SbX a -> SbZ a -> SbU a -> SbP a -> a -> (SbR a, SbO a) +sbDae :: forall a . Floating a => SbX a -> SbX a -> SbZ a -> SbU a -> SbP a -> None a -> a -> (SbR a, SbO a) sbDae (SbX gamma' p' v') (SbX gamma _ v@(V2 vx vz)) @@ -109,6 +112,7 @@ (SbU omega alpha) _ _ + _ = (residual, outputs) where residual :: SbR a @@ -170,13 +174,14 @@ , bcP0 :: V2 a } deriving (Functor, Generic, Generic1, Show) -bc :: Num a => SbX a -> SbX a -> None a -> SbP a -> a -> SbBc a +bc :: Num a => SbX a -> SbX a -> None a -> SbP a -> None a -> a -> SbBc a bc (SbX gamma0 p0@(V2 _ pz0) (V2 vx0 vz0)) (SbX gammaF (V2 _ pzF) (V2 vxF vzF)) _ _ _ + _ = SbBc { bcPeriodicGamma = gamma0 + gammaF , bcPeriodicPz = pz0 - pzF @@ -185,11 +190,11 @@ , bcP0 = p0 } -mayer :: Floating a => a -> SbX a -> SbX a -> None a -> SbP a -> a -mayer tf _ (SbX _ (V2 pxF _) _) _ _ = - pxF / tf +mayer :: Floating a => a -> SbX a -> SbX a -> None a -> SbP a -> None a -> a +mayer tf _ (SbX _ (V2 pxF _) _) _ _ _ = - pxF / tf -lagrange :: Floating a => SbX a -> SbZ a -> SbU a -> SbP a -> SbO a -> a -> a -> a -lagrange _ _ (SbU omega alpha) _ _ _ _ = 1e-3*omega*omega + 1e-3*alpha*alpha +lagrange :: Floating a => SbX a -> SbZ a -> SbU a -> SbP a -> None a -> SbO a -> a -> a -> a +lagrange _ _ (SbU omega alpha) _ _ _ _ _ = 1e-3*omega*omega + 1e-3*alpha*alpha ubnd :: SbU (Maybe Double, Maybe Double) ubnd = SbU @@ -206,36 +211,46 @@ (Just (-100), Just 100) (Just (-100), Just 100)) -pathc :: t -> t1 -> t2 -> t3 -> t4 -> t5 -> None a -pathc _ _ _ _ _ _ = None +pathc :: t -> t1 -> t2 -> t3 -> t4 -> t5 -> t6 -> None a +pathc _ _ _ _ _ _ _ = None ocp :: OcpPhase' SailboatOcp -ocp = OcpPhase { ocpMayer = mayer - , ocpLagrange = lagrange - , ocpQuadratures = \_ _ _ _ _ _ _ -> None - , ocpDae = sbDae - , ocpBc = bc - , ocpPathC = pathc - , ocpPathCBnds = None - , ocpBcBnds = fill (Just 0, Just 0) - , ocpXbnd = xbnd - , ocpUbnd = ubnd - , ocpZbnd = SbZ - , ocpPbnd = SbP - , ocpTbnd = (Just 1, Just 50) - , ocpObjScale = Nothing - , ocpTScale = Nothing - , ocpXScale = Nothing - , ocpZScale = Nothing - , ocpUScale = Nothing - , ocpPScale = Nothing - , ocpResidualScale = Nothing - , ocpBcScale = Nothing - , ocpPathCScale = Nothing - } +ocp = + OcpPhase + { ocpMayer = mayer + , ocpLagrange = lagrange + , ocpQuadratures = \_ _ _ _ _ _ _ _ -> None + , ocpQuadratureOutputs = \_ _ _ _ _ _ _ _ -> None + , ocpDae = sbDae + , ocpBc = bc + , ocpPlotOutputs = \_ _ _ _ _ _ _ _ _ _ _ -> None + , ocpPathC = pathc + , ocpObjScale = Nothing + , ocpTScale = Nothing + , ocpXScale = Nothing + , ocpZScale = Nothing + , ocpUScale = Nothing + , ocpPScale = Nothing + , ocpResidualScale = Nothing + , ocpBcScale = Nothing + , ocpPathCScale = Nothing + } +ocpInputs :: OcpPhaseInputs' SailboatOcp +ocpInputs = + OcpPhaseInputs + { ocpPathCBnds = None + , ocpBcBnds = fill (Just 0, Just 0) + , ocpXbnd = xbnd + , ocpUbnd = ubnd + , ocpZbnd = SbZ + , ocpPbnd = SbP + , ocpTbnd = (Just 1, Just 50) + , ocpFixedP = None + } + urlDynoPlot :: String urlDynoPlot = "tcp://127.0.0.1:5563" @@ -277,16 +292,16 @@ main :: IO () main = do - cp <- makeCollProblem Legendre ocp (cat initialGuess) + cp <- makeCollProblem Legendre ocp ocpInputs (cat initialGuess) let nlp = cpNlp cp ZMQ.withContext $ \context -> withPublisher context urlDynoPlot $ \sendDynoPlotMsg -> do -- withPublisher context urlOptTelem $ \sendOptTelemMsg -> do let meta = toMeta (cpMetaProxy cp) - callback :: J (CollTraj' SailboatOcp NCollStages CollDeg) (Vector Double) -> IO Bool - callback traj = do - plotPoints <- cpPlotPoints cp traj + callback :: J (CollTraj' SailboatOcp NCollStages CollDeg) (Vector Double) -> b -> IO Bool + callback traj _ = do + plotPoints <- cpPlotPoints cp traj (catJV None) -- dynoplot let dynoPlotMsg = encodeSerial (plotPoints, meta) sendDynoPlotMsg "glider" dynoPlotMsg diff --git a/examples/Sofa/Common.hs b/examples/Sofa/Common.hs --- a/examples/Sofa/Common.hs +++ b/examples/Sofa/Common.hs @@ -1,7 +1,7 @@ {-# OPTIONS_GHC -Wall #-} -{-# Language DeriveFunctor #-} -{-# Language DeriveGeneric #-} -{-# Language PolyKinds #-} +{-# LANGUAGE DeriveFunctor #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE PolyKinds #-} module Sofa.Common ( SofaMessage(..) diff --git a/examples/SofaExpando.hs b/examples/SofaExpando.hs --- a/examples/SofaExpando.hs +++ b/examples/SofaExpando.hs @@ -1,10 +1,10 @@ -- | How big of a sofa can we get around a corner? {-# OPTIONS_GHC -Wall #-} -{-# Language DeriveFunctor #-} -{-# Language DeriveGeneric #-} -{-# Language ScopedTypeVariables #-} -{-# Language DataKinds #-} +{-# LANGUAGE DeriveFunctor #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE DataKinds #-} module Main where diff --git a/examples/SofaVisualizer.hs b/examples/SofaVisualizer.hs --- a/examples/SofaVisualizer.hs +++ b/examples/SofaVisualizer.hs @@ -1,5 +1,5 @@ {-# OPTIONS_GHC -Wall #-} -{-# Language CPP #-} +{-# LANGUAGE CPP #-} module Main ( main ) where @@ -114,15 +114,18 @@ ) where walls = VisObjects - [ Line [ V3 (-4) 1 0 + [ Line Nothing + [ V3 (-4) 1 0 , V3 1 1 0 , V3 1 (-4) 0 ] (makeColor 1 1 1 1) - , Line [ V3 (-4) 2 0 + , Line Nothing + [ V3 (-4) 2 0 , V3 2 2 0 , V3 2 (-4) 0 ] (makeColor 1 1 1 1) - , Line [ V3 (-4) 0 0 + , Line Nothing + [ V3 (-4) 0 0 , V3 0 0 0 , V3 0 (-4) 0 ] (makeColor 1 1 1 1) @@ -131,13 +134,13 @@ axes = Axes (0.5, 15) npoints = length points nsteps = length stages - shape0 = Line' $ + shape0 = Line' Nothing $ zipWith (\(Point x y) c -> ((V3 x y 0) - (V3 2 2 0), c)) (points ++ [head points]) (colors (npoints + 1)) drawOne :: [Point Double] -> Double -> Color -> VisObject Double drawOne ps@(p0:_) z = - Line + Line Nothing (map (\(Point x y) -> (V3 x y z)) (ps ++ [p0])) drawOne _ _ = const (VisObjects []) diff --git a/examples/Spring.hs b/examples/Spring.hs --- a/examples/Spring.hs +++ b/examples/Spring.hs @@ -1,8 +1,8 @@ {-# OPTIONS_GHC -Wall #-} -{-# Language TypeFamilies #-} -{-# Language DeriveFunctor #-} -{-# Language DeriveGeneric #-} -{-# Language DataKinds #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE DeriveFunctor #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE DataKinds #-} module Main ( SpringX(..), SpringU(..), main ) where @@ -13,6 +13,7 @@ import Accessors ( Lookup ) import Dyno.View.View ( J, jfill ) +import Dyno.View.JV ( catJV ) import Dyno.Nlp ( Bounds ) import Dyno.Ocp import Dyno.Vectorize ( Vectorize, None(..), fill ) @@ -29,17 +30,12 @@ OcpPhase { ocpMayer = mayer , ocpLagrange = lagrange - , ocpQuadratures = \_ _ _ _ _ _ _ -> None + , ocpQuadratures = \_ _ _ _ _ _ _ _ -> None + , ocpQuadratureOutputs = \_ _ _ _ _ _ _ _ -> None , ocpDae = dae , ocpBc = bc , ocpPathC = pathC - , ocpBcBnds = bcBnds - , ocpPathCBnds = pathCBnds - , ocpXbnd = fill (Nothing, Nothing) - , ocpZbnd = fill (Nothing, Nothing) - , ocpUbnd = fill (Nothing, Nothing) - , ocpPbnd = fill (Nothing, Nothing) - , ocpTbnd = (Just 4, Just 4) + , ocpPlotOutputs = \_ _ _ _ _ _ _ _ _ _ _ -> None , ocpObjScale = Nothing , ocpTScale = Nothing , ocpXScale = Nothing @@ -51,6 +47,19 @@ , ocpPathCScale = Nothing } +springOcpInputs :: OcpPhaseInputs' SpringOcp +springOcpInputs = + OcpPhaseInputs + { ocpBcBnds = bcBnds + , ocpPathCBnds = pathCBnds + , ocpXbnd = fill (Nothing, Nothing) + , ocpZbnd = fill (Nothing, Nothing) + , ocpUbnd = fill (Nothing, Nothing) + , ocpPbnd = fill (Nothing, Nothing) + , ocpTbnd = (Just 4, Just 4) + , ocpFixedP = None + } + data SpringOcp type instance X SpringOcp = SpringX type instance O SpringOcp = SpringO @@ -61,6 +70,9 @@ type instance P SpringOcp = None type instance Z SpringOcp = None type instance Q SpringOcp = None +type instance QO SpringOcp = None +type instance FP SpringOcp = None +type instance PO SpringOcp = None data SpringX a = SpringX @@ -94,9 +106,9 @@ instance Lookup a => Lookup (SpringPathC a) dae :: Floating a - => SpringX a -> SpringX a -> None a -> SpringU a -> None a -> a + => SpringX a -> SpringX a -> None a -> SpringU a -> None a -> None a -> a -> (SpringX a, SpringO a) -dae (SpringX p' v') (SpringX p v) _ (SpringU u) _ t = +dae (SpringX p' v') (SpringX p v) _ (SpringU u) _ _ t = (residual, outputs) where residual = SpringX (p' - v) (v' - force) @@ -109,8 +121,8 @@ force = u - k * p - b * v + 0.1 * sin t obj = p**2 + v**2 + u**2 -bc :: SpringX a -> SpringX a -> None a -> None a -> a -> SpringBc a -bc x0 xf _ _ _ = SpringBc x0 xf +bc :: SpringX a -> SpringX a -> None a -> None a -> None a -> a -> SpringBc a +bc x0 xf _ _ _ _ = SpringBc x0 xf bcBnds :: SpringBc Bounds bcBnds = @@ -119,18 +131,18 @@ , bcXF = SpringX (Just 1, Just 1) (Just 0, Just 0) } -mayer :: Floating a => a -> SpringX a -> SpringX a -> None a -> None a -> a -mayer endTime _ (SpringX pf vf) _ _ = (pf**2 + vf**2 + endTime/1000) +mayer :: Floating a => a -> SpringX a -> SpringX a -> None a -> None a -> None a -> a +mayer endTime _ (SpringX pf vf) _ _ _ = (pf**2 + vf**2 + endTime/1000) -pathC :: Floating a => SpringX a -> None a -> SpringU a -> None a -> SpringO a -> a -> SpringPathC a -pathC (SpringX _ v) _ (SpringU u) _ _ time = +pathC :: Floating a => SpringX a -> None a -> SpringU a -> None a -> None a -> SpringO a -> a -> SpringPathC a +pathC (SpringX _ v) _ (SpringU u) _ _ _ time = SpringPathC (v**2 + u**2 - time/100) pathCBnds :: SpringPathC Bounds pathCBnds = SpringPathC (Nothing, Just 4) -lagrange :: Fractional a => SpringX a -> None a -> SpringU a -> None a -> SpringO a -> a -> a -> a -lagrange _ _ _ _ (SpringO force obj) _ _ = obj + force*force*1e-4 +lagrange :: Fractional a => SpringX a -> None a -> SpringU a -> None a -> None a -> SpringO a -> a -> a -> a +lagrange _ _ _ _ _ (SpringO force obj) _ _ = obj + force*force*1e-4 solver :: Solver solver = ipoptSolver { options = [("expand", Opt True)] } @@ -145,12 +157,12 @@ main = withCallback $ \send -> do - cp <- makeCollProblem Legendre springOcp guess + cp <- makeCollProblem Legendre springOcp springOcpInputs guess let nlp = cpNlp cp meta = toMeta (cpMetaProxy cp) - cb' traj = do - plotPoints <- cpPlotPoints cp traj + cb' traj _ = do + plotPoints <- cpPlotPoints cp traj (catJV None) send (plotPoints, meta) _ <- solveNlp solver nlp (Just cb') diff --git a/examples/ToyOcp.hs b/examples/ToyOcp.hs new file mode 100644 --- /dev/null +++ b/examples/ToyOcp.hs @@ -0,0 +1,46 @@ +{-# OPTIONS_GHC -Wall #-} +{-# LANGUAGE DeriveFunctor #-} +{-# LANGUAGE DeriveGeneric #-} + +module Main where + +import GHC.Generics ( Generic1 ) + +import Dyno.Vectorize +import Dyno.SimpleOcp + +-- state +data X a = X { xTheta :: a, xOmega :: a } + deriving (Functor, Generic1, Show) +instance Vectorize X + +-- controls +data U a = U { uTorque :: a } + deriving (Functor, Generic1, Show) +instance Vectorize U + +pendOde :: Floating a => X a -> U a -> X a +pendOde (X theta omega) (U torque) = X thetaDot omegaDot + where + thetaDot = omega + omegaDot = torque + 9.8 * sin theta + +ocp :: SimpleOcp X U +ocp = + SimpleOcp + { ode = pendOde + , objective = \(X _ omega) (U torque) -> omega * omega + torque * torque + , xBounds = X (-pi, pi) (-5, 5) + , uBounds = U (-50, 50) + , xInitial = X {xTheta = pi/2, xOmega = 0} + , xFinal = X {xTheta = 0, xOmega = 0} + , endTime = 1 + , initialGuess = \t -> X ((1-t) * pi/2) (pi/1) + } + +main :: IO () +main = do + result <- solveOcp ocp + case result of + Left msg -> putStrLn $ "failed with " ++ msg + Right xus -> print xus diff --git a/examples/Vec.hs b/examples/Vec.hs --- a/examples/Vec.hs +++ b/examples/Vec.hs @@ -2,11 +2,11 @@ -- Don't forget to import DataKinds/PolyKinds ! {-# OPTIONS_GHC -Wall #-} -{-# Language ScopedTypeVariables #-} -{-# Language DeriveFunctor #-} -{-# Language DeriveGeneric #-} -{-# Language DataKinds #-} -{-# Language PolyKinds #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE DeriveFunctor #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE PolyKinds #-} module Main where @@ -40,7 +40,7 @@ -- you do know the length at compile time knownLength :: (Num a, Show a) => Vec 4 (Params a) -knownLength = mkVec unknownLength +knownLength = devectorize unknownLength -- do something on type-safe vec data doSomething :: (Dim n, Num a) => Vec n (Params a) -> a diff --git a/examples/beginner/SimpleQp.hs b/examples/beginner/SimpleQp.hs --- a/examples/beginner/SimpleQp.hs +++ b/examples/beginner/SimpleQp.hs @@ -1,6 +1,6 @@ {-# OPTIONS_GHC -Wall #-} -{-# Language DeriveFunctor #-} -{-# Language DeriveGeneric #-} +{-# LANGUAGE DeriveFunctor #-} +{-# LANGUAGE DeriveGeneric #-} module Main where diff --git a/src/Dyno/AutoScaling.hs b/src/Dyno/AutoScaling.hs --- a/src/Dyno/AutoScaling.hs +++ b/src/Dyno/AutoScaling.hs @@ -1,7 +1,7 @@ {-# OPTIONS_GHC -Wall #-} -{-# Language ScopedTypeVariables #-} -{-# Language DeriveFunctor #-} -{-# Language DeriveGeneric #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE DeriveFunctor #-} +{-# LANGUAGE DeriveGeneric #-} module Dyno.AutoScaling ( scalingNlp diff --git a/src/Dyno/DirectCollocation/ActiveConstraints.hs b/src/Dyno/DirectCollocation/ActiveConstraints.hs new file mode 100644 --- /dev/null +++ b/src/Dyno/DirectCollocation/ActiveConstraints.hs @@ -0,0 +1,213 @@ +{-# OPTIONS_GHC -Wall #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE DeriveFunctor #-} +{-# LANGUAGE DeriveFoldable #-} +{-# LANGUAGE DeriveTraversable #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE PolyKinds #-} + +module Dyno.DirectCollocation.ActiveConstraints + ( ActiveConstraints(..) + , Active(..) + , getActiveConstraints + , flattenActiveConstraints + , summarizeActiveConstraints + , matlabActiveConstraints + , pythonActiveConstraints + ) where + +import GHC.Generics ( Generic ) + +import Control.Applicative +import Data.List ( intercalate ) +import Data.Maybe ( catMaybes ) +import qualified Data.Foldable as F +import qualified Data.Traversable as T +import Data.Vector ( Vector ) +import Text.Printf ( printf ) + +import Dyno.DirectCollocation.Types +import Dyno.Ocp ( OcpPhase(..), OcpPhaseInputs(..) ) +import Dyno.Nlp ( Bounds ) +import Dyno.Vectorize ( Vectorize, Id(..) ) +import Dyno.View.View ( View(..), J ) +import Dyno.View.JV ( JV, splitJV ) +import Dyno.View.JVec ( unJVec ) +import Dyno.TypeVecs ( Dim ) + +import Accessors ( Lookup, Getter(..), flatten', accessors ) + +data Active a = Active { activeLower :: a, activeUpper :: a } + deriving (Functor, F.Foldable, T.Traversable, Generic) +instance Lookup a => Lookup (Active a) + +data ActiveConstraints x z u p h c a = + ActiveConstraints + { xBounds :: x a + , zBounds :: z a + , uBounds :: u a + , pBounds :: p a + , pathConstraintBounds :: h a + , boundaryConditionBounds :: c a + , endTimeBounds :: a + } deriving (Functor, F.Foldable, T.Traversable, Generic) +instance ( Lookup (x a), Lookup (z a), Lookup (u a), Lookup (p a) + , Lookup (h a), Lookup (c a), Lookup a + ) => Lookup (ActiveConstraints x z u p h c a) + +summarizeActiveConstraints :: + ( Functor x, Functor z, Functor u, Functor p, Functor h, Functor c + , Lookup (x Int) + , Lookup (z Int) + , Lookup (u Int) + , Lookup (p Int) + , Lookup (h Int) + , Lookup (c Int) + ) => ActiveConstraints x z u p h c (Active Int) -> String +summarizeActiveConstraints activeCons = + unlines $ catMaybes $ map report $ flattenActiveConstraints activeCons + where + report (_, Active 0 0) = Nothing + report (name, Active lb ub) = + Just $ printf "% 4d lower, % 4d upper (%s)" lb ub (intercalate "." name) + +matlabActiveConstraints :: + ( Functor x, Functor z, Functor u, Functor p, Functor h, Functor c + , Lookup (x Int) + , Lookup (z Int) + , Lookup (u Int) + , Lookup (p Int) + , Lookup (h Int) + , Lookup (c Int) + ) => ActiveConstraints x z u p h c (Active Int) -> String +matlabActiveConstraints activeCons = "{" ++ intercalate "; " cons ++ "}" + where + cons = map report $ flattenActiveConstraints activeCons + report (name, Active lb ub) = printf "'%s', %d, %d" (intercalate "." name) lb ub + +pythonActiveConstraints :: + ( Functor x, Functor z, Functor u, Functor p, Functor h, Functor c + , Lookup (x Int) + , Lookup (z Int) + , Lookup (u Int) + , Lookup (p Int) + , Lookup (h Int) + , Lookup (c Int) + ) => ActiveConstraints x z u p h c (Active Int) -> String +pythonActiveConstraints activeCons = "[" ++ intercalate ", " cons ++ "]" + where + cons = map report $ flattenActiveConstraints activeCons + report (name, Active lb ub) = printf "('%s', %d, %d)" (intercalate "." name) lb ub + +flattenActiveConstraints :: + forall x z u p h c . + ( Functor x, Functor z, Functor u, Functor p, Functor h, Functor c + , Lookup (x Int) + , Lookup (z Int) + , Lookup (u Int) + , Lookup (p Int) + , Lookup (h Int) + , Lookup (c Int) + ) => ActiveConstraints x z u p h c (Active Int) -> [([String], Active Int)] +flattenActiveConstraints activeCons = map report $ flatten' $ accessors lbs + where + report (name, GetInt get, _) = (name, Active (get lbs) (get ubs)) + report (name, _, _) = + error $ "the 'impossible' happened, flattenActiveConstraints got a non-int getter " ++ show name + lbs = fmap activeLower activeCons + ubs = fmap activeUpper activeCons + +getActiveConstraints :: + forall x z u p h c n deg r fp o q qo po + . ( Vectorize x, Vectorize z, Vectorize u, Vectorize p, Vectorize h, Vectorize c, Vectorize r + , Applicative x, Applicative z, Applicative u, Applicative p, Applicative h, Applicative c + , Dim n, Dim deg + ) + => (J (CollTraj x z u p n deg) (Vector Double) + -> J (JV fp) (Vector Double) + -> IO (J (CollOcpConstraints x r c h n deg) (Vector Double)) + ) + -> OcpPhase x z u p r o c h q qo po fp + -> Double + -> J (CollTraj x z u p n deg) (Vector Double) + -> J (JV fp) (Vector Double) + -> OcpPhaseInputs x z u p c h fp + -> IO (ActiveConstraints x z u p h c (Active Int)) +getActiveConstraints evalConstraints ocp eps x p inputs = do + g <- evalConstraints x p + return $ whatsActive eps (split x) (split g) inputs ocp + +whatsActive :: + forall x z u p h c n deg r fp o q qo po + . ( Vectorize x, Vectorize z, Vectorize u, Vectorize p, Vectorize h, Vectorize c + , Applicative x, Applicative z, Applicative u, Applicative p, Applicative h, Applicative c + , Dim n, Dim deg + ) + => Double + -> CollTraj x z u p n deg (Vector Double) + -> CollOcpConstraints x r c h n deg (Vector Double) + -> OcpPhaseInputs x z u p c h fp + -> OcpPhase x z u p r o c h q qo po fp + -> ActiveConstraints x z u p h c (Active Int) +whatsActive userEps traj@(CollTraj tf p _ _) g inputs ocp = + ActiveConstraints + { xBounds = countEmAll $ map (isActive (ocpXScale ocp) (ocpXbnd inputs)) xs + , zBounds = countEmAll $ map (isActive (ocpZScale ocp) (ocpZbnd inputs)) zs + , uBounds = countEmAll $ map (isActive (ocpUScale ocp) (ocpUbnd inputs)) us + , pBounds = isActive (ocpPScale ocp) (ocpPbnd inputs) (splitJV p) + , pathConstraintBounds = countEmAll $ map (isActive (ocpPathCScale ocp) (ocpPathCBnds inputs)) pathC + , boundaryConditionBounds = isActive (ocpBcScale ocp) (ocpBcBnds inputs) bc + , endTimeBounds = scalarIsActive userEps (ocpTScale ocp) (ocpTbnd inputs) (unId (splitJV tf)) + } + where + countEmAll :: forall f . Applicative f => [f (Active Int)] -> f (Active Int) + countEmAll counts = liftA2 Active lowers uppers + where + lowers :: f Int + lowers = fmap sum $ T.sequenceA $ map (fmap activeLower) counts + + uppers :: f Int + uppers = fmap sum $ T.sequenceA $ map (fmap activeUpper) counts + + pathC :: [h Double] + pathC = concatMap (map splitJV . F.toList . unJVec . split) $ F.toList $ unJVec $ split (coPathC g) + (xs', zs', us') = getXzus'' traj + xs = concatMap F.toList (F.toList xs') + zs = concatMap F.toList (F.toList zs') + us = concatMap F.toList (F.toList us') + + isActive :: Applicative f => Maybe (f Double) -> f Bounds -> f Double -> f (Active Int) + isActive scale bnds val = (scalarIsActive userEps) <$> T.sequenceA scale <*> bnds <*> val + + bc :: c Double + bc = splitJV (coBc g) + + +scalarIsActive :: Double -> Maybe Double -> Bounds -> Double -> Active Int +scalarIsActive _ _ (Nothing, Nothing) _ = Active 0 0 +scalarIsActive userEps scale bnd@(Just lb, Nothing) x + | eps >= x - lb = Active 1 0 -- lower active + | otherwise = Active 0 0 + where + eps = toEps userEps scale bnd +scalarIsActive userEps scale bnd@(Nothing, Just ub) x + | eps >= ub - x = Active 0 1 -- upper active + | otherwise = Active 0 0 + where + eps = toEps userEps scale bnd +scalarIsActive userEps scale bnd@(Just lb, Just ub) x + | lb == ub = Active 0 0 -- don't report equality constraints, duh + | eps >= x - lb = Active 1 0 -- lower active + | eps >= ub - x = Active 0 1 -- upper active + | otherwise = Active 0 0 + where + eps = toEps userEps scale bnd + + +toEps :: Double -> Maybe Double -> (Maybe Double, Maybe Double) -> Double +toEps userEps (Just scale) _ = scale * userEps +toEps userEps Nothing (Just lb, Just ub) = scale * userEps + where + scale = ub - lb +toEps userEps Nothing _ = userEps diff --git a/src/Dyno/DirectCollocation/Dynamic.hs b/src/Dyno/DirectCollocation/Dynamic.hs --- a/src/Dyno/DirectCollocation/Dynamic.hs +++ b/src/Dyno/DirectCollocation/Dynamic.hs @@ -1,10 +1,11 @@ {-# OPTIONS_GHC -Wall #-} -{-# Language ScopedTypeVariables #-} -{-# Language DeriveGeneric #-} -{-# Language PolyKinds #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE PolyKinds #-} +-- todo(greg): rename to PlotPoints or something module Dyno.DirectCollocation.Dynamic - ( DynPlotPoints + ( DynPlotPoints(..) , CollTrajMeta(..) , addCollocationChannel , MetaProxy(..) @@ -35,7 +36,7 @@ import Dyno.View.Unsafe.View ( unJ, unJ' ) -import Dyno.Vectorize ( Vectorize, Id(..), fill ) +import Dyno.Vectorize ( Vectorize(..), Id(..), fill ) import Dyno.View.JV ( JV, splitJV ) import Dyno.View.View ( View(..), J ) import Dyno.View.JVec ( JVec(..) ) @@ -44,7 +45,6 @@ import Dyno.DirectCollocation.Types import Dyno.DirectCollocation.Quadratures ( QuadratureRoots, mkTaus ) - addCollocationChannel :: String -> (((DynPlotPoints Double, CollTrajMeta) -> IO ()) -> IO ()) -> Plotter () addCollocationChannel name action = addChannel name sameMeta toSignalTree action @@ -60,14 +60,7 @@ sameMeta :: (DynPlotPoints Double, CollTrajMeta) -> (DynPlotPoints Double, CollTrajMeta) -> Bool -sameMeta (_,ctm0) (_,ctm1) = - and [ ctmX ctm0 == ctmX ctm1 - , ctmZ ctm0 == ctmZ ctm1 - , ctmU ctm0 == ctmU ctm1 - , ctmP ctm0 == ctmP ctm1 - , ctmO ctm0 == ctmO ctm1 - , ctmQ ctm0 == ctmQ ctm1 - ] +sameMeta (_,ctm0) (_,ctm1) = ctm0 == ctm1 data DynPlotPoints a = DynPlotPoints (Vector (Vector (a, Vector a))) @@ -76,6 +69,9 @@ (Vector (Vector (a, Vector a))) (Vector (Vector (a, Vector a))) (Vector (Vector (a, Vector a))) + (Vector (Vector (a, Vector a))) + (Vector (Vector (a, Vector a))) + (Vector (Vector (a, Vector a))) deriving Generic @@ -85,27 +81,26 @@ catDynPlotPoints :: V.Vector (DynPlotPoints a) -> DynPlotPoints a catDynPlotPoints pps = DynPlotPoints - (V.concatMap (\(DynPlotPoints x _ _ _ _ _) -> x) pps) - (V.concatMap (\(DynPlotPoints _ x _ _ _ _) -> x) pps) - (V.concatMap (\(DynPlotPoints _ _ x _ _ _) -> x) pps) - (V.concatMap (\(DynPlotPoints _ _ _ x _ _) -> x) pps) - (V.concatMap (\(DynPlotPoints _ _ _ _ x _) -> x) pps) - (V.concatMap (\(DynPlotPoints _ _ _ _ _ x) -> x) pps) + (V.concatMap (\(DynPlotPoints x _ _ _ _ _ _ _ _) -> x) pps) + (V.concatMap (\(DynPlotPoints _ x _ _ _ _ _ _ _) -> x) pps) + (V.concatMap (\(DynPlotPoints _ _ x _ _ _ _ _ _) -> x) pps) + (V.concatMap (\(DynPlotPoints _ _ _ x _ _ _ _ _) -> x) pps) + (V.concatMap (\(DynPlotPoints _ _ _ _ x _ _ _ _) -> x) pps) + (V.concatMap (\(DynPlotPoints _ _ _ _ _ x _ _ _) -> x) pps) + (V.concatMap (\(DynPlotPoints _ _ _ _ _ _ x _ _) -> x) pps) + (V.concatMap (\(DynPlotPoints _ _ _ _ _ _ _ x _) -> x) pps) + (V.concatMap (\(DynPlotPoints _ _ _ _ _ _ _ _ x) -> x) pps) dynPlotPoints :: - forall x z u p h o n deg a . + forall x z u p h o q qo po n deg a . ( Dim n, Dim deg, Real a, Fractional a, Show a - , Vectorize x, Vectorize z, Vectorize u, Vectorize o, Vectorize p, Vectorize h + , Vectorize x, Vectorize z, Vectorize u, Vectorize o, Vectorize p, Vectorize h, Vectorize q + , Vectorize po, Vectorize qo ) => QuadratureRoots -> CollTraj x z u p n deg (Vector a) - -> Vec n ( Vec deg ( J (JV o) (Vector a) - , J (JV x) (Vector a) - , J (JV h) (Vector a) - ) - , J (JV x) (Vector a) - ) + -> Vec n (StageOutputs x o h q qo po deg a) -> DynPlotPoints a dynPlotPoints quadratureRoots (CollTraj tf' _ stages' xf) outputs -- if degree is one, each arc will be 1 point and won't get drawn @@ -113,8 +108,11 @@ -- https://github.com/ghorn/Plot-ho-matic/issues/10 -- https://github.com/timbod7/haskell-chart/issues/81 | reflectDim (Proxy :: Proxy deg) == 1 = - DynPlotPoints xss' (singleArc zss) (singleArc uss) (singleArc oss) (singleArc xdss) (singleArc hss) - | otherwise = DynPlotPoints xss' zss uss oss xdss hss + DynPlotPoints xss (singleArc zss) (singleArc uss) (singleArc oss) (singleArc xdss) + (singleArc hss) (singleArc poss) (singleArc qss) (singleArc qdss) + | otherwise = DynPlotPoints xss zss uss oss xdss hss poss (singleArc qss) qdss + -- draw quadrature states as a single line since they are differentiable + -- and gaps are closed by construction where singleArc :: Vector (Vector b) -> Vector (Vector b) singleArc = V.singleton . V.concat . V.toList @@ -129,18 +127,21 @@ stages :: Vec n (CollStage (JV x) (JV z) (JV u) deg (Vector a)) stages = fmap split (unJVec (split stages')) - xss' = xss `V.snoc` (V.singleton (tf, unJ xf)) + xss = xss' `V.snoc` (V.singleton (tf, unJ xf)) + -- assumes initial time is 0 + qss = V.singleton (0, vectorize (fill 0 :: Quadratures q qo a)) `V.cons` qss' - xss,zss,uss,oss,xdss,hss :: Vector (Vector (a, Vector a)) - (xss,zss,uss,oss,xdss,hss) = V.unzip6 xzuoxdhs + xss',zss,uss,oss,poss,xdss,hss :: Vector (Vector (a, Vector a)) + (xss',zss,uss,oss,xdss,hss,poss,qss',qdss) = unzip9 xzuoxdhs -- todo: check this final time tf'' against expected tf (_tf'', xzuoxdhs) = T.mapAccumL f 0 $ V.zip (TV.unVec stages) (TV.unVec outputs) + -- todo(greg): should take the times from toCallbacks, not recalculate f :: a -> ( CollStage (JV x) (JV z) (JV u) deg (Vector a) - , (Vec deg (J (JV o) (Vector a), J (JV x) (Vector a), J (JV h) (Vector a)), J (JV x) (Vector a)) + , StageOutputs x o h q qo po deg a ) -> ( a , ( V.Vector (a, V.Vector a) @@ -149,32 +150,57 @@ , V.Vector (a, V.Vector a) , V.Vector (a, V.Vector a) , V.Vector (a, V.Vector a) + , V.Vector (a, V.Vector a) + , V.Vector (a, V.Vector a) + , V.Vector (a, V.Vector a) ) ) - f t0 (CollStage x0 xzus', (xdohs, xnext)) = (tnext, (xs,zs,us,os,xds,hs)) + f t0 (CollStage x0 xzus', stageOutputs) = (tnext, (xs,zs,us,os,xds,hs,pos,qs,qds)) where tnext = t0 + h xzus0 = fmap split (unJVec (split xzus')) :: Vec deg (CollPoint (JV x) (JV z) (JV u) (Vector a)) xs :: V.Vector (a, V.Vector a) - xs = (t0, unJ x0) `V.cons` xs' `V.snoc` (tnext,unJ xnext) + xs = (t0, unJ x0) `V.cons` xs' `V.snoc` (tnext, unJ (soXNext stageOutputs)) - xs',zs,us,os,xds,hs :: Vector (a, Vector a) - (xs',zs,us,os,xds,hs) = V.unzip6 $ TV.unVec $ TV.tvzipWith3 g xzus0 xdohs taus + qs :: V.Vector (a, V.Vector a) + qs = qs' `V.snoc` (tnext, vectorize (soQNext stageOutputs)) - g (CollPoint x z u) (o,x',pathc) tau = + xs',zs,us,os,xds,hs,pos,qs',qds :: Vector (a, Vector a) + (xs',zs,us,os,xds,hs,pos,qs',qds) = + unzip9 $ TV.unVec $ TV.tvzipWith3 g xzus0 (soVec stageOutputs) taus + + g :: CollPoint (JV x) (JV z) (JV u) (Vector a) + -> ( J (JV o) (Vector a), J (JV x) (Vector a), J (JV h) (Vector a) + , J (JV po) (Vector a) + , Quadratures q qo a, Quadratures q qo a + ) + -> a + -> ( (a, V.Vector a) + , (a, V.Vector a) + , (a, V.Vector a) + , (a, V.Vector a) + , (a, V.Vector a) + , (a, V.Vector a) + , (a, V.Vector a) + , (a, V.Vector a) + , (a, V.Vector a) + ) + g (CollPoint x z u) (o,x',pathc,po,q,q') tau = ( (t,unJ' "x" x) , (t,unJ' "z" z) , (t,unJ' "u" u) , (t,unJ' "o" o) , (t,unJ' "x'" x') , (t,unJ' "h" pathc) + , (t,unJ' "po" po) + , (t,vectorize q) + , (t,vectorize q') ) where t = t0 + h*tau - data NameTree = NameTreeNode (String,String) [(String,NameTree)] | NameTreeLeaf Int deriving (Show, Eq, Generic) @@ -188,6 +214,7 @@ , ctmO :: NameTree , ctmQ :: NameTree , ctmH :: NameTree + , ctmPo :: NameTree } deriving (Eq, Generic, Show) instance Binary CollTrajMeta instance Serialize CollTrajMeta @@ -205,14 +232,17 @@ type MetaTree a = Tree.Forest (String, String, Maybe ((DynPlotPoints a, CollTrajMeta) -> [[(a,a)]])) forestFromMeta :: CollTrajMeta -> MetaTree Double -forestFromMeta meta = [xTree,zTree,uTree,oTree,xdTree,hTree] +forestFromMeta meta = [xTree,zTree,uTree,oTree,xdTree,hTree,poTree,qTree,qdTree] where - xTree = blah (\(DynPlotPoints x _ _ _ _ _) -> x) "differential states" (ctmX meta) - zTree = blah (\(DynPlotPoints _ z _ _ _ _) -> z) "algebraic variables" (ctmZ meta) - uTree = blah (\(DynPlotPoints _ _ u _ _ _) -> u) "controls" (ctmU meta) - oTree = blah (\(DynPlotPoints _ _ _ o _ _) -> o) "outputs" (ctmO meta) - xdTree = blah (\(DynPlotPoints _ _ _ _ xd _) -> xd) "diff state derivatives" (ctmX meta) - hTree = blah (\(DynPlotPoints _ _ _ _ _ h) -> h) "path constraints" (ctmH meta) + xTree = blah (\(DynPlotPoints x _ _ _ _ _ _ _ _ ) -> x) "differential states" (ctmX meta) + zTree = blah (\(DynPlotPoints _ z _ _ _ _ _ _ _ ) -> z) "algebraic variables" (ctmZ meta) + uTree = blah (\(DynPlotPoints _ _ u _ _ _ _ _ _ ) -> u) "controls" (ctmU meta) + oTree = blah (\(DynPlotPoints _ _ _ o _ _ _ _ _ ) -> o) "outputs" (ctmO meta) + xdTree = blah (\(DynPlotPoints _ _ _ _ xd _ _ _ _ ) -> xd) "diff state derivatives" (ctmX meta) + hTree = blah (\(DynPlotPoints _ _ _ _ _ h _ _ _ ) -> h) "path constraints" (ctmH meta) + poTree = blah (\(DynPlotPoints _ _ _ _ _ _ po _ _ ) -> po) "quadrature outputs" (ctmPo meta) + qTree = blah (\(DynPlotPoints _ _ _ _ _ _ _ q _ ) -> q) "quadrature states" (ctmQ meta) + qdTree = blah (\(DynPlotPoints _ _ _ _ _ _ _ _ qd) -> qd) "ddt(quadrature states)" (ctmQ meta) blah :: forall f c t . (Functor f, F.Foldable f) @@ -226,15 +256,15 @@ woo = F.toList . fmap (F.toList . fmap (\(t,x) -> (t, x V.! k))) -data MetaProxy x z u p o q h = MetaProxy +data MetaProxy x z u p o q qo po h = MetaProxy -toMeta :: forall x z u p o q h . +toMeta :: forall x z u p o q qo po h . ( Lookup (x ()), Lookup (z ()), Lookup (u ()), Lookup (p ()), Lookup (o ()), Lookup (q ()) - , Lookup (h ()) + , Lookup (h ()), Lookup (po ()), Lookup (qo ()) , Vectorize x, Vectorize z, Vectorize u, Vectorize p, Vectorize o, Vectorize q - , Vectorize h + , Vectorize h, Vectorize po, Vectorize qo ) - => MetaProxy x z u p o q h -> CollTrajMeta + => MetaProxy x z u p o q qo po h -> CollTrajMeta toMeta _ = CollTrajMeta { ctmX = namesFromAccTree $ accessors (fill () :: x ()) @@ -242,6 +272,32 @@ , ctmU = namesFromAccTree $ accessors (fill () :: u ()) , ctmP = namesFromAccTree $ accessors (fill () :: p ()) , ctmO = namesFromAccTree $ accessors (fill () :: o ()) - , ctmQ = namesFromAccTree $ accessors (fill () :: q ()) + , ctmQ = namesFromAccTree $ accessors (fill () :: Quadratures q qo ()) , ctmH = namesFromAccTree $ accessors (fill () :: h ()) + , ctmPo = namesFromAccTree $ accessors (fill () :: po ()) } + +--unzip8 :: Vector (a, b, c, d, e, f, g, h) +-- -> (Vector a, Vector b, Vector c, Vector d, Vector e, Vector f, Vector g, Vector h) +--{-# INLINE unzip8 #-} +--unzip8 xs = (V.map (\(a, _, _, _, _, _, _, _) -> a) xs, +-- V.map (\(_, b, _, _, _, _, _, _) -> b) xs, +-- V.map (\(_, _, c, _, _, _, _, _) -> c) xs, +-- V.map (\(_, _, _, d, _, _, _, _) -> d) xs, +-- V.map (\(_, _, _, _, e, _, _, _) -> e) xs, +-- V.map (\(_, _, _, _, _, f, _, _) -> f) xs, +-- V.map (\(_, _, _, _, _, _, g, _) -> g) xs, +-- V.map (\(_, _, _, _, _, _, _, h) -> h) xs) + +unzip9 :: Vector (a, b, c, d, e, f, g, h, i) + -> (Vector a, Vector b, Vector c, Vector d, Vector e, Vector f, Vector g, Vector h, Vector i) +{-# INLINE unzip9 #-} +unzip9 xs = (V.map (\(a, _, _, _, _, _, _, _, _) -> a) xs, + V.map (\(_, b, _, _, _, _, _, _, _) -> b) xs, + V.map (\(_, _, c, _, _, _, _, _, _) -> c) xs, + V.map (\(_, _, _, d, _, _, _, _, _) -> d) xs, + V.map (\(_, _, _, _, e, _, _, _, _) -> e) xs, + V.map (\(_, _, _, _, _, f, _, _, _) -> f) xs, + V.map (\(_, _, _, _, _, _, g, _, _) -> g) xs, + V.map (\(_, _, _, _, _, _, _, h, _) -> h) xs, + V.map (\(_, _, _, _, _, _, _, _, i) -> i) xs) diff --git a/src/Dyno/DirectCollocation/Export.hs b/src/Dyno/DirectCollocation/Export.hs --- a/src/Dyno/DirectCollocation/Export.hs +++ b/src/Dyno/DirectCollocation/Export.hs @@ -1,53 +1,122 @@ {-# OPTIONS_GHC -Wall #-} -{-# Language ScopedTypeVariables #-} -{-# Language PolyKinds #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE PolyKinds #-} module Dyno.DirectCollocation.Export - ( toMatlab + ( Export(..) + , ExportConfig(..) + , exportTraj + , exportTraj' + -- * matlab specific + , matlabParam + , matlabTraj + -- * python specific + , PythonExporter + , runPythonExporter + , pythonParam + , pythonTraj + , write ) where +import Control.Monad ( unless ) +import Data.List ( unzip6, intercalate ) import Data.Proxy ( Proxy(..) ) import Linear.V ( Dim(..) ) import Data.Vector ( Vector ) -import qualified Data.Vector as V import qualified Data.Foldable as F - -import Accessors ( Lookup, flatten, accessors ) +import Control.Monad.State.Lazy ( State ) +import qualified Control.Monad.State.Lazy as State +import qualified Data.Set as S -import Dyno.View.Unsafe.View ( unJ ) +import Accessors ( Lookup, Getter(..), flatten, flatten', accessors ) import Dyno.Nlp ( NlpOut(..) ) import Dyno.TypeVecs ( Vec ) -import Dyno.Vectorize ( Vectorize, Id(..), fill ) +import Dyno.Vectorize ( Vectorize, Id(..), None(..), fill ) import Dyno.View.View ( View(..) ) -import Dyno.View.JV ( JV, splitJV ) -import Dyno.View.JVec ( JVec(..) ) +import Dyno.View.JV ( splitJV, catJV ) import Dyno.DirectCollocation.Formulate ( CollProblem(..) ) -import Dyno.DirectCollocation.Types ( CollTraj(..), CollStage(..), CollPoint(..) ) +import Dyno.DirectCollocation.Types ( CollTraj(..), CollOcpConstraints(..) + , StageOutputs(..), Quadratures(..) + , getXzus''' + ) import Dyno.DirectCollocation.Quadratures ( timesFromTaus ) -toMatlab :: - forall x z u p r o c h q n deg lol +data ExportConfig = + ExportConfig + { ecMatlabVariableName :: String + , ecPythonVariableName :: String + } + +data Export = + Export + { exportMatlab :: String + , exportPython :: String + } + +exportTraj :: + forall x z u p fp r o c h q qo po n deg . ( Lookup (x Double), Vectorize x , Lookup (z Double), Vectorize z , Lookup (u Double), Vectorize u , Lookup (o Double), Vectorize o , Lookup (p Double), Vectorize p + , Lookup (c Double), Vectorize c + , Vectorize r + , Lookup (fp Double), Vectorize fp , Lookup (h Double), Vectorize h + , Lookup (q Double), Vectorize q + , Lookup (po Double), Vectorize po + , Lookup (qo Double), Vectorize qo , Dim n, Dim deg ) - => CollProblem x z u p r o c h q n deg - -> NlpOut (CollTraj x z u p n deg) lol (Vector Double) - -> IO String -toMatlab cp nlpOut = do - let ct@(CollTraj tf' p' stages' xf) = split (xOpt nlpOut) + => ExportConfig + -> CollProblem x z u p r o c h q qo po fp n deg + -> fp Double + -> NlpOut (CollTraj x z u p n deg) (CollOcpConstraints x r c h n deg) (Vector Double) + -> IO Export +exportTraj = exportTraj' (Nothing :: Maybe ([String], None Double)) + + +-- | this version takes optional user data +exportTraj' :: + forall x z u p fp r o c h q qo po n deg e + . ( Lookup (x Double), Vectorize x + , Lookup (z Double), Vectorize z + , Lookup (u Double), Vectorize u + , Lookup (o Double), Vectorize o + , Lookup (p Double), Vectorize p + , Lookup (c Double), Vectorize c + , Vectorize r + , Lookup (fp Double), Vectorize fp + , Lookup (h Double), Vectorize h + , Lookup (q Double), Vectorize q + , Lookup (po Double), Vectorize po + , Lookup (qo Double), Vectorize qo + , Dim n, Dim deg + , Lookup (e Double), Vectorize e + ) + => Maybe ([String], e Double) + -> ExportConfig + -> CollProblem x z u p r o c h q qo po fp n deg + -> fp Double + -> NlpOut (CollTraj x z u p n deg) (CollOcpConstraints x r c h n deg) (Vector Double) + -> IO Export +exportTraj' mextra exportConfig cp fp nlpOut = do + let matlabRetName = ecMatlabVariableName exportConfig + pyRetName = ecPythonVariableName exportConfig + + let ct@(CollTraj tf' p' _ _) = split (xOpt nlpOut) CollTraj lagTf' lagP' _ _ = split (lambdaXOpt nlpOut) + lagBc' = coBc $ split (lambdaGOpt nlpOut) - outs <- cpOutputs cp (cat ct) + (_, outs, finalQuads) <- cpHellaOutputs cp (cat ct) (catJV fp) + let _ = outs :: Vec n (StageOutputs x o h q qo po deg Double) + _ = finalQuads :: Quadratures q qo Double let taus :: Vec deg Double taus = cpTaus cp - tf = V.head (unJ tf') + Id tf = splitJV tf' n = reflectDim (Proxy :: Proxy n) @@ -59,61 +128,195 @@ xTimes = concatMap (\(t0,ts) -> t0 : F.toList ts) (F.toList times) ++ [tf] zuoTimes = concatMap (\(_,ts) -> F.toList ts) (F.toList times) - stages :: [CollStage (JV x) (JV z) (JV u) deg (Vector Double)] - stages = map split $ F.toList $ unJVec $ split stages' + xss :: Vec n (x Double, Vec deg (x Double)) + xf :: x Double + zss :: Vec n (Vec deg (z Double)) + uss :: Vec n (Vec deg (u Double)) + ((xss,xf), zss, uss) = getXzus''' ct + fullXs :: [x Double] + fullXs = concatMap (\(x0, xs') -> x0 : F.toList xs') (F.toList xss) ++ [xf] + xs :: [x Double] - xs = concatMap getXs stages ++ [splitJV xf] + xs = concatMap (F.toList . snd) (F.toList xss) zs :: [z Double] - zs = concatMap getZs stages + zs = concatMap F.toList (F.toList zss) us :: [u Double] - us = concatMap getUs stages + us = concatMap F.toList (F.toList uss) os :: [o Double] xdots :: [x Double] hs :: [h Double] - (os, xdots, hs) = unzip3 $ F.concatMap (F.toList . fst) outs -- drop the interpolated value + -- drop the interpolated value + os = map splitJV os' + xdots = map splitJV xdots' + hs = map splitJV hs' + pos = map splitJV pos' + (os', xdots', hs', pos', _, _) = unzip6 $ F.concatMap (F.toList . soVec) outs + qsFull :: [Quadratures q qo Double] + qsFull = fill 0 : F.concatMap toQFull outs + where + toQFull :: StageOutputs x o h q qo po deg Double -> [Quadratures q qo Double] + toQFull stageOutputs = (map (\(_,_,_,_,qs',_) -> qs') (F.toList (soVec stageOutputs))) ++ [soQNext stageOutputs] - getXs (CollStage x0 xzus) = splitJV x0 : map (getX . split) (F.toList (unJVec (split xzus))) - getZs (CollStage _ xzus) = map (getZ . split) (F.toList (unJVec (split xzus))) - getUs (CollStage _ xzus) = map (getU . split) (F.toList (unJVec (split xzus))) + qs :: [Quadratures q qo Double] + qs = F.concatMap toQ outs + where + toQ :: StageOutputs x o h q qo po deg Double -> [Quadratures q qo Double] + toQ stageOutputs = map (\(_,_,_,_,qs',_) -> qs') (F.toList (soVec stageOutputs)) - getX :: CollPoint (JV x) (JV z) (JV u) (Vector Double) -> x Double - getX (CollPoint x _ _) = splitJV x + toQd :: StageOutputs x o h q qo po deg Double -> [Quadratures q qo Double] + toQd stageOutputs = (map (\(_,_,_,_,_,qd) -> qd) (F.toList (soVec stageOutputs))) + qds :: [Quadratures q qo Double] + qds = F.concatMap toQd outs - getZ :: CollPoint (JV x) (JV z) (JV u) (Vector Double) -> z Double - getZ (CollPoint _ z _) = splitJV z + matlabOut :: String + matlabOut = unlines $ + matlabTraj (matlabRetName ++ ".diffStatesFull") fullXs ++ + matlabTraj (matlabRetName ++ ".diffStates") xs ++ + matlabTraj (matlabRetName ++ ".diffStateDerivs") xdots ++ + matlabTraj (matlabRetName ++ ".algVars") zs ++ + matlabTraj (matlabRetName ++ ".controls") us ++ + matlabTraj (matlabRetName ++ ".outputs") os ++ + matlabTraj (matlabRetName ++ ".pathConstraints") hs ++ + matlabTraj (matlabRetName ++ ".plotOutputs") pos ++ + matlabTraj (matlabRetName ++ ".quadratureStatesFull") qsFull ++ + matlabTraj (matlabRetName ++ ".quadratureStates") qs ++ + matlabTraj (matlabRetName ++ ".quadratureStateDerivs") qds ++ + matlabParam (matlabRetName ++ ".params") (splitJV p') ++ + ( case mextra of + Nothing -> [] + Just (names,extra) -> matlabParam (intercalate "." (matlabRetName : names)) extra + ) ++ + matlabParam (matlabRetName ++ ".lagrangeMultipliers.params") (splitJV lagP') ++ + matlabParam (matlabRetName ++ ".lagrangeMultipliers.bc") (splitJV lagBc') ++ + matlabParam (matlabRetName ++ ".finalQuadratureStates") finalQuads ++ + [ matlabRetName ++ ".lagrangeMultipliers.T = " ++ show (unId (splitJV lagTf')) ++ ";" + , "" + , matlabRetName ++ ".tx = " ++ show xTimes ++ ";" + , matlabRetName ++ ".tzuo = " ++ show zuoTimes ++ ";" + , matlabRetName ++ ".T = " ++ show tf ++ ";" + , matlabRetName ++ ".N = " ++ show n ++ ";" + , matlabRetName ++ ".deg = " ++ show (reflectDim (Proxy :: Proxy deg)) ++ ";" + , matlabRetName ++ ".collocationRoots = '" ++ show (cpRoots cp) ++ "';" + ] - getU :: CollPoint (JV x) (JV z) (JV u) (Vector Double) -> u Double - getU (CollPoint _ _ u) = splitJV u + runRet :: State PythonExporter () + runRet = do + write "import numpy" + write "" + write $ pyRetName ++ " = {}" + pythonTraj pyRetName ["diffStatesFull"] fullXs + pythonTraj pyRetName ["diffStates"] xs + pythonTraj pyRetName ["diffStateDerivs"] xdots + pythonTraj pyRetName ["algVars"] zs + pythonTraj pyRetName ["controls"] us + pythonTraj pyRetName ["outputs"] os + pythonTraj pyRetName ["pathConstraints"] hs + pythonTraj pyRetName ["plotOutputs"] pos + pythonTraj pyRetName ["quadratureStatesFull"] qsFull + pythonTraj pyRetName ["quadratureStates"] qs + pythonTraj pyRetName ["quadratureStateDerivs"] qds + pythonParam pyRetName ["params"] (splitJV p') + case mextra of + Nothing -> return () + Just (names,extra) -> pythonParam pyRetName names extra + pythonParam pyRetName ["lagrangeMultipliers","params"] (splitJV lagP') + pythonParam pyRetName ["lagrangeMultipliers","bc"] (splitJV lagBc') + pythonParam pyRetName ["finalQuadratureStates"] finalQuads + putVal pyRetName ["lagrangeMultipliers","T"] (show (unId (splitJV lagTf'))) + write "" + putVal pyRetName ["tx"] (npArray (show xTimes)) + putVal pyRetName ["tzuo"] (npArray (show zuoTimes)) + putVal pyRetName ["T"] (show tf) + putVal pyRetName ["N"] (show n) + putVal pyRetName ["deg"] (show (reflectDim (Proxy :: Proxy deg))) + putVal pyRetName ["collocationRoots"] ("'" ++ show (cpRoots cp) ++ "'") - at :: (Vectorize xzu, Lookup (xzu Double)) => [(String, xzu Double -> Double)] - at = flatten $ accessors (fill 0) + return $ Export + { exportMatlab = matlabOut + , exportPython = unlines (runPythonExporter runRet) + } - woo :: String -> [xzu Double] -> String -> (xzu Double -> Double) -> String - woo topName xzus name get = topName ++ "." ++ name ++ " = " ++ show (map get xzus) ++ ";" +runPythonExporter :: State PythonExporter () -> [String] +runPythonExporter action = reverse pythonOut + where + PythonExporter (_, pythonOut) = State.execState action (PythonExporter (S.empty, [])) - wooP :: String -> p Double -> String -> (p Double -> Double) -> String - wooP topName p name get = topName ++ "." ++ name ++ " = " ++ show (get p) ++ ";" +npArray :: String -> String +npArray str = "numpy.array(" ++ str ++ ")" - ret :: String - ret = init $ unlines $ - map (uncurry (woo "ret.diffStates" xs)) at ++ - map (uncurry (woo "ret.diffStateDerivs" xdots)) at ++ - map (uncurry (woo "ret.algVars" zs)) at ++ - map (uncurry (woo "ret.controls" us)) at ++ - map (uncurry (woo "ret.outputs" os)) at ++ - map (uncurry (woo "ret.pathConstraints" hs)) at ++ - map (uncurry (wooP "ret.params" (splitJV p'))) at ++ - map (uncurry (wooP "ret.lagrangeMultipliers.params" (splitJV lagP'))) at ++ - [ "ret.lagrangeMultipliers.T = " ++ show (unId (splitJV lagTf')) - , "" - , "ret.tx = " ++ show xTimes - , "ret.tzuo = " ++ show zuoTimes - , "ret.N = " ++ show n - , "ret.deg = " ++ show (reflectDim (Proxy :: Proxy deg)) - , "ret.collocationRoots = '" ++ show (cpRoots cp) ++ "'" - ] - return ret +toDub :: Getter (xzu Double) -> xzu Double -> Double +toDub (GetDouble f) = f +toDub (GetFloat f) = realToFrac . f +toDub (GetInt f) = realToFrac . f +toDub (GetBool f) = fromIntegral . fromEnum . f +toDub (GetString _) = const (read "NaN") +toDub GetSorry = const (read "NaN") + + +pythonParam :: forall p . (Vectorize p, Lookup (p Double)) + => String -> [String] -> p Double -> State PythonExporter () +pythonParam pyRetName topNames p = mapM_ pyParam at' + where + pyParam :: ([String], (p Double -> Double)) -> State PythonExporter () + pyParam (name, get) = putVal pyRetName (topNames ++ name) (show (get p)) + + at' :: [([String], p Double -> Double)] + at' = map (\(fn,g,_) -> (fn, toDub g)) $ flatten' $ accessors (fill (0 :: Double)) + +pythonTraj :: forall x . (Vectorize x, Lookup (x Double)) + => String -> [String] -> [x Double] -> State PythonExporter () +pythonTraj pyRetName topNames xs = mapM_ pyArray at' + where + pyArray :: ([String], (x Double -> Double)) -> State PythonExporter () + pyArray (name, get) = putVal pyRetName (topNames ++ name) (npArray (show (map get xs))) + + at' :: [([String], x Double -> Double)] + at' = map (\(fn,g,_) -> (fn, toDub g)) $ flatten' $ accessors (fill (0 :: Double)) + + +matlabParam :: forall p . (Vectorize p, Lookup (p Double)) => String -> p Double -> [String] +matlabParam topName p = map (uncurry mlParam) at + where + mlParam :: String -> (p Double -> Double) -> String + mlParam name get = topName ++ "." ++ name ++ " = " ++ show (get p) ++ ";" + + at :: [(String, p Double -> Double)] + at = map (\(fn,g,_) -> (fn, toDub g)) $ flatten $ accessors (fill (0 :: Double)) + +matlabTraj :: forall x . (Vectorize x, Lookup (x Double)) => String -> [x Double] -> [String] +matlabTraj topName xs = map (uncurry mlArray) at + where + mlArray :: String -> (x Double -> Double) -> String + mlArray name get = + topName ++ "." ++ name ++ " = " ++ show (map get xs) ++ ";" + + at :: [(String, x Double -> Double)] + at = map (\(fn,g,_) -> (fn, toDub g)) $ flatten $ accessors (fill (0 :: Double)) + +data PythonExporter = PythonExporter (S.Set [String], [String]) + +pyname :: String -> [String] -> String +pyname topName xs = topName ++ concatMap (\x -> "['" ++ x ++ "']") xs + +putNameIfMissing :: String -> [String] -> State PythonExporter () +putNameIfMissing _ [] = return () +putNameIfMissing topName name = do + PythonExporter (set0, _) <- State.get + unless (S.member name set0) $ do + putNameIfMissing topName (init name) + PythonExporter (set1, out1) <- State.get + State.put $ PythonExporter (S.insert name set1, (pyname topName name ++ " = {}") : out1) + +write :: String -> State PythonExporter () +write str = do + PythonExporter (set0, outs0) <- State.get + State.put $ PythonExporter (set0, str:outs0) + +putVal :: String -> [String] -> String -> State PythonExporter () +putVal topName name val = do + putNameIfMissing topName name + write (pyname topName name ++ " = " ++ val) diff --git a/src/Dyno/DirectCollocation/Formulate.hs b/src/Dyno/DirectCollocation/Formulate.hs --- a/src/Dyno/DirectCollocation/Formulate.hs +++ b/src/Dyno/DirectCollocation/Formulate.hs @@ -1,1061 +1,1064 @@ {-# OPTIONS_GHC -Wall #-} -{-# Language TypeFamilies #-} -{-# Language DeriveGeneric #-} -{-# Language ScopedTypeVariables #-} -{-# Language TypeOperators #-} -{-# Language FlexibleContexts #-} -{-# Language PolyKinds #-} - -module Dyno.DirectCollocation.Formulate - ( CovTraj(..) - , CollProblem(..) - , makeCollProblem - , CollCovProblem(..) - , makeCollCovProblem - , mkTaus - , makeGuess - , makeGuessSim - ) where - -import GHC.Generics ( Generic ) - -import Data.Maybe ( fromMaybe ) -import Data.Proxy ( Proxy(..) ) -import Data.Vector ( Vector ) -import qualified Data.Vector as V -import qualified Data.Foldable as F -import qualified Data.Traversable as T -import qualified Data.Packed.Matrix as Mat -import qualified Numeric.LinearAlgebra.Algorithms as LA -import Linear.Matrix hiding ( trace ) -import Linear.V - -import Casadi.DMatrix ( DMatrix ) -import Casadi.MX ( MX ) - -import Dyno.SXElement ( sxCatJV, sxSplitJV ) -import Dyno.View.View ( View(..), J, jfill, JTuple(..), JNone(..), v2d, d2v ) -import qualified Dyno.View.M as M -import Dyno.View.Cov ( Cov ) -import Dyno.View.JV ( JV, splitJV, catJV, catJV' ) -import Dyno.View.HList ( (:*:)(..) ) -import Dyno.View.Fun -import Dyno.View.JVec( JVec(..), jreplicate ) -import Dyno.View.Viewable ( Viewable ) -import Dyno.View.Scheme ( Scheme ) -import Dyno.Vectorize ( Vectorize(..), Id(..), fill, vlength, vzipWith ) -import Dyno.TypeVecs ( Vec ) -import qualified Dyno.TypeVecs as TV -import Dyno.LagrangePolynomials ( lagrangeDerivCoeffs ) -import Dyno.Nlp ( Nlp(..), Bounds ) -import Dyno.Ocp - -import Dyno.DirectCollocation.Types -import Dyno.DirectCollocation.Dynamic ( MetaProxy(..), DynPlotPoints, dynPlotPoints ) -import Dyno.DirectCollocation.Quadratures ( QuadratureRoots(..), mkTaus, interpolate, timesFromTaus ) -import Dyno.DirectCollocation.Robust - -data CollProblem x z u p r o c h q n deg = - CollProblem - { cpNlp :: Nlp (CollTraj x z u p n deg) - JNone - (CollOcpConstraints x r c h n deg) MX - , cpOcp :: OcpPhase x z u p r o c h q - , cpPlotPoints :: J (CollTraj x z u p n deg) (Vector Double) - -> IO (DynPlotPoints Double) - , cpHellaOutputs :: J (CollTraj x z u p n deg) (Vector Double) - -> IO ( DynPlotPoints Double - , Vec n ( Vec deg ( J (JV o) (Vector Double) - , J (JV x) (Vector Double) - , J (JV h) (Vector Double) - ) - , J (JV x) (Vector Double) - ) - ) - , cpOutputs :: J (CollTraj x z u p n deg) (Vector Double) - -> IO (Vec n ( Vec deg ( o Double - , x Double - , h Double - ) - , x Double - ) - ) - , cpTaus :: Vec deg Double - , cpRoots :: QuadratureRoots - , cpEvalQuadratures :: Vec n (Vec deg Double) -> Double -> IO Double - , cpMetaProxy :: MetaProxy x z u p o q h - } - -makeCollProblem :: - forall x z u p r o c h q deg n . - ( Dim deg, Dim n - , Vectorize x, Vectorize p, Vectorize u, Vectorize z - , Vectorize r, Vectorize o, Vectorize h, Vectorize c, Vectorize q - ) - => QuadratureRoots -> OcpPhase x z u p r o c h q - -> J (CollTraj x z u p n deg) (Vector Double) - -> IO (CollProblem x z u p r o c h q n deg) -makeCollProblem roots ocp guess = do - let -- the collocation points - taus :: Vec deg Double - taus = mkTaus roots - - n = reflectDim (Proxy :: Proxy n) - - -- coefficients for getting xdot by lagrange interpolating polynomials - cijs :: Vec (TV.Succ deg) (Vec (TV.Succ deg) Double) - cijs = lagrangeDerivCoeffs (0 TV.<| taus) - - interpolate' :: (J (JV x) :*: J (JVec deg (JV x))) MX -> J (JV x) MX - interpolate' (x0 :*: xs) = case roots of - Legendre -> interpolate taus x0 (unJVec (split xs)) - Radau -> TV.tvlast $ unJVec $ split xs - - interpolateq' :: (J (JV q) :*: J (JVec deg (JV q))) MX -> J (JV q) MX - interpolateq' (q0 :*: qs) = case roots of - Legendre -> interpolate taus q0 (unJVec (split qs)) - Radau -> TV.tvlast $ unJVec $ split qs - - interpolateScalar' :: (J (JV Id) :*: J (JVec deg (JV Id))) MX -> J (JV Id) MX - interpolateScalar' (x0 :*: xs) = case roots of - Legendre -> interpolate taus x0 (unJVec (split xs)) - Radau -> TV.tvlast $ unJVec $ split xs - - dynamicsFunction (t :*: parm :*: x' :*: collPoint) = (sxCatJV r) :*: (sxCatJV o) - where - CollPoint x z u = split collPoint - (r,o) = ocpDae ocp - (sxSplitJV x') (sxSplitJV x) (sxSplitJV z) (sxSplitJV u) - (sxSplitJV parm) (unId (sxSplitJV t)) - - interpolateFun <- toMXFun "interpolate (JV x)" interpolate' >>= expandMXFun - interpolateQFun <- toMXFun "interpolate (JV q)" interpolateq' >>= expandMXFun - interpolateScalarFun <- toMXFun "interpolate (JV Id)" interpolateScalar' >>= expandMXFun - let callInterpolateScalar :: J (JV Id) MX -> Vec deg (J (JV Id) MX) -> J (JV Id) MX - callInterpolateScalar x0 xs = call interpolateScalarFun (x0 :*: cat (JVec xs)) - - callInterpolate :: J (JV x) MX -> Vec deg (J (JV x) MX) -> J (JV x) MX - callInterpolate x0 xs = call interpolateFun (x0 :*: cat (JVec xs)) - - callInterpolateQ :: J (JV q) MX -> Vec deg (J (JV q) MX) -> J (JV q) MX - callInterpolateQ q0 qs = call interpolateQFun (q0 :*: cat (JVec qs)) - - bcFun <- toSXFun "bc" $ \(x0:*:x1:*:x2:*:x3:*:x4) -> sxCatJV $ ocpBc ocp (sxSplitJV x0) (sxSplitJV x1) (sxSplitJV x2) (sxSplitJV x3) (unId (sxSplitJV x4)) - mayerFun <- toSXFun "mayer" $ \(x0:*:x1:*:x2:*:x3:*:x4) -> - sxCatJV $ Id $ ocpMayer ocp (unId (sxSplitJV x0)) (sxSplitJV x1) (sxSplitJV x2) (sxSplitJV x3) (sxSplitJV x4) - - lagrangeFun <- toSXFun "lagrange" $ \(x0:*:x1:*:x2:*:x3:*:x4:*:x5:*:x6) -> - sxCatJV $ Id $ ocpLagrange ocp (sxSplitJV x0) (sxSplitJV x1) (sxSplitJV x2) (sxSplitJV x3) (sxSplitJV x4) (unId (sxSplitJV x5)) (unId (sxSplitJV x6)) - lagQuadFun <- toMXFun "lagrange quadratures" $ evaluateQuadraturesFunction lagrangeFun callInterpolateScalar cijs n - callLagQuadFun <- fmap call (expandMXFun lagQuadFun) -- necessary to discard unused outputs - - quadratureDotFun <- toSXFun "quadrature derivative" $ \(x0:*:x1:*:x2:*:x3:*:x4:*:x5:*:x6) -> - sxCatJV $ ocpQuadratures ocp (sxSplitJV x0) (sxSplitJV x1) (sxSplitJV x2) (sxSplitJV x3) (sxSplitJV x4) (unId (sxSplitJV x5)) (unId (sxSplitJV x6)) - quadFun <- toMXFun "quadratures" $ evaluateQuadraturesFunction quadratureDotFun callInterpolateQ cijs n - callQuadFun <- fmap call (expandMXFun quadFun) -- necessary to discard unused outputs - - genericQuadraturesFun <- toMXFun "generic quadratures" $ genericQuadraturesFunction callInterpolateScalar cijs n - - dynFun <- toSXFun "dynamics" dynamicsFunction - - pathConFun <- toSXFun "pathConstraints" $ pathConFunction $ - \x0 x1 x2 x3 x4 x5 -> sxCatJV $ ocpPathC ocp (sxSplitJV x0) (sxSplitJV x1) (sxSplitJV x2) (sxSplitJV x3) (sxSplitJV x4) (unId (sxSplitJV x5)) - pathStageConFun <- toMXFun "pathStageCon" (pathStageConstraints pathConFun) - - dynStageConFun <- toMXFun "dynamicsStageCon" (dynStageConstraints callInterpolate cijs dynFun) - - stageFun <- toMXFun "stageFunction" $ stageFunction pathStageConFun (call dynStageConFun) --- let callStageFun = call stageFun - callStageFun <- fmap call (expandMXFun stageFun) - - outputFun <- toMXFun "stageOutputs" $ outputFunction callInterpolate cijs taus dynFun - - -- prepare callbacks - let f :: J (JV o) DMatrix -> J (JV x) DMatrix -> J (JV h) DMatrix - -> (J (JV o) (Vector Double), J (JV x) (Vector Double), J (JV h) (Vector Double)) - f o' x' h' = (d2v o', d2v x', d2v h') - - callOutputFun :: J (JV p) (Vector Double) - -> J (JV Id) (Vector Double) - -> J (CollStage (JV x) (JV z) (JV u) deg) (Vector Double) - -> J (JV Id) (Vector Double) - -> IO ( Vec deg ( J (JV o) (Vector Double) - , J (JV x) (Vector Double) - , J (JV h) (Vector Double) - ) - , J (JV x) (Vector Double) - ) - callOutputFun p h stage k = do - let p' = v2d p - (_ :*: xdot :*: out :*: xnext) <- - eval outputFun $ (v2d stage) :*: p' :*: (v2d h) :*: (v2d k) - - let stageTimes :: Vec deg (J (JV Id) DMatrix) - stageTimes = fmap (\tau -> t0 + realToFrac tau * h') taus - where - t0 = h' * v2d k - h' = v2d h - CollStage _ collPoints = split stage - hs <- eval pathStageConFun $ p' :*: (cat (JVec stageTimes)) :*: out :*: (v2d collPoints) - - let outs0 = unJVec (split out) :: Vec deg (J (JV o) DMatrix) - xdots0 = unJVec (split xdot) :: Vec deg (J (JV x) DMatrix) - hs0 = unJVec (split hs) :: Vec deg (J (JV h) DMatrix) - return (TV.tvzipWith3 f outs0 xdots0 hs0, d2v xnext) - - mapOutputFun :: J (CollTraj x z u p n deg) (Vector Double) - -> IO (Vec n ( Vec deg ( J (JV o) (Vector Double) - , J (JV x) (Vector Double) - , J (JV h) (Vector Double) - ) - , J (JV x) (Vector Double) - ) - ) - mapOutputFun ct = do - let CollTraj tf p stages _ = split ct - h = catJV $ Id (tf' / fromIntegral n) - where - Id tf' = splitJV tf - - vstages = unJVec (split stages) - :: Vec n (J (CollStage (JV x) (JV z) (JV u) deg) (Vector Double)) - ks :: Vec n (J (JV Id) (Vector Double)) - ks = TV.mkVec' $ map (catJV . Id . realToFrac) (take n [(0::Int)..]) - - T.sequence $ TV.tvzipWith (callOutputFun p h) vstages ks - - getHellaOutputs :: - J (CollTraj x z u p n deg) (Vector Double) - -> IO ( DynPlotPoints Double - , Vec n ( Vec deg ( J (JV o) (Vector Double) - , J (JV x) (Vector Double) - , J (JV h) (Vector Double) - ) - , J (JV x) (Vector Double) - ) - ) - getHellaOutputs traj = do - outputs <- mapOutputFun traj - return (dynPlotPoints roots (split traj) outputs, outputs) - - getPlotPoints :: J (CollTraj x z u p n deg) (Vector Double) - -> IO (DynPlotPoints Double) - getPlotPoints traj = fmap fst $ getHellaOutputs traj - - getOutputs :: J (CollTraj x z u p n deg) (Vector Double) - -> IO (Vec n (Vec deg (o Double, x Double, h Double), x Double)) - getOutputs traj = do - outputs <- mapOutputFun traj - let devec :: Vec deg (J (JV o) (Vector Double), J (JV x) (Vector Double), J (JV h) (Vector Double)) - -> Vec deg (o Double, x Double, h Double) - devec = fmap (\(x,y,z) -> (splitJV x, splitJV y, splitJV z)) - return $ fmap (\(x,y) -> (devec x, splitJV y)) outputs - - let nlp :: Nlp (CollTraj x z u p n deg) JNone (CollOcpConstraints x r c h n deg) MX - nlp = Nlp { - nlpFG = - getFg taus - (bcFun :: SXFun (J (JV x) :*: J (JV x) :*: J (JV q) :*: J (JV p) :*: J (JV Id)) (J (JV c))) - (mayerFun :: SXFun (J (JV Id) :*: (J (JV x) :*: (J (JV x)) :*: (J (JV q)) :*: (J (JV p)))) (J (JV Id))) - (callLagQuadFun :: (J (JV p) :*: J (JVec deg (CollPoint (JV x) (JV z) (JV u))) :*: J (JVec deg (JV o)) :*: J (JV Id) :*: J (JVec deg (JV Id))) MX - -> J (JV Id) MX) - (callQuadFun :: (J (JV p) :*: J (JVec deg (CollPoint (JV x) (JV z) (JV u))) :*: J (JVec deg (JV o)) :*: J (JV Id) :*: J (JVec deg (JV Id))) MX - -> J (JV q) MX) - (callStageFun :: (J (JV Id) :*: J (JV p) :*: J (JVec deg (JV Id)) :*: J (JV x) :*: J (JVec deg (JTuple (JV x) (JV z))) :*: J (JVec deg (JV u))) MX - -> (J (JVec deg (JV r)) :*: J (JVec deg (JV o)) :*: J (JVec deg (JV h)) :*: J (JV x)) MX) - , nlpBX = cat $ fillCollTraj' - (fill (Nothing, Nothing)) - (ocpXbnd ocp) - (ocpZbnd ocp) - (ocpUbnd ocp) - (ocpPbnd ocp) - (ocpTbnd ocp) - , nlpBG = cat (getBg ocp) - , nlpX0 = guess :: J (CollTraj x z u p n deg) (Vector Double) - , nlpP = cat JNone - , nlpLamX0 = Nothing - , nlpLamG0 = Nothing - , nlpScaleF = ocpObjScale ocp - , nlpScaleX = Just $ cat $ fillCollTraj - (fromMaybe (fill 1) (ocpXScale ocp)) - (fromMaybe (fill 1) (ocpZScale ocp)) - (fromMaybe (fill 1) (ocpUScale ocp)) - (fromMaybe (fill 1) (ocpPScale ocp)) - (fromMaybe 1 (ocpTScale ocp)) - - , nlpScaleG = Just $ cat $ fillCollConstraints - (fromMaybe (fill 1) (ocpXScale ocp)) - (fromMaybe (fill 1) (ocpResidualScale ocp)) - (fromMaybe (fill 1) (ocpBcScale ocp)) - (fromMaybe (fill 1) (ocpPathCScale ocp)) - } - evalQuadratures :: Vec n (Vec deg Double) -> Double -> IO Double - evalQuadratures qs' tf' = do - let d2d :: Double -> J (JV Id) DMatrix - d2d = realToFrac - qs :: Vec n (J (JVec deg (JV Id)) DMatrix) - qs = fmap (cat . JVec . fmap d2d) qs' - tf :: J (JV Id) DMatrix - tf = realToFrac tf' - evalq :: J (JVec deg (JV Id)) DMatrix -> IO (J (JV Id) DMatrix) - evalq q = eval genericQuadraturesFun (q :*: tf) - stageIntegrals' <- T.mapM evalq qs :: IO (Vec n (J (JV Id) DMatrix)) - let stageIntegrals = fmap (unId . splitJV . d2v) stageIntegrals' :: Vec n Double - return (F.sum stageIntegrals) - - return $ CollProblem { cpNlp = nlp - , cpOcp = ocp - , cpPlotPoints = getPlotPoints - , cpHellaOutputs = getHellaOutputs - , cpOutputs = getOutputs - , cpTaus = taus - , cpRoots = roots - , cpEvalQuadratures = evalQuadratures - , cpMetaProxy = MetaProxy - } - - -data CollCovProblem ocp n deg sx sw sh shr sc = - CollCovProblem - { ccpNlp :: Nlp - (CollTrajCov sx ocp n deg) - JNone - (CollOcpCovConstraints ocp n deg sh shr sc) MX - , ccpPlotPoints :: J (CollTrajCov sx ocp n deg) (Vector Double) -> IO (DynPlotPoints Double) - , ccpOutputs :: - J (CollTrajCov sx ocp n deg) (Vector Double) - -> IO ( Vec n (Vec deg (O ocp Double, X ocp Double, H ocp Double), X ocp Double) - , Vec n (J (Cov (JV sx)) (Vector Double)) - , J (Cov (JV sx)) (Vector Double) - ) - , ccpSensitivities :: MXFun - (J (CollTraj' ocp n deg)) - (CovarianceSensitivities (JV sx) (JV sw) n) - , ccpCovariances :: MXFun - (J (CollTrajCov sx ocp n deg)) (J (CovTraj sx n)) - , ccpRoots :: QuadratureRoots - } - -makeCollCovProblem :: - forall ocp x z u p r o c h q sx sz sw sr sh shr sc deg n . - ( Dim deg, Dim n, Vectorize x, Vectorize p, Vectorize u, Vectorize z - , Vectorize sr, Vectorize sw, Vectorize sz, Vectorize sx - , Vectorize r, Vectorize o, Vectorize h, Vectorize c, Vectorize q - , View sh, Vectorize shr, View sc - , x ~ X ocp - , q ~ Q ocp - , h ~ H ocp - , c ~ C ocp - , o ~ O ocp - , r ~ R ocp - , p ~ P ocp - , u ~ U ocp - , z ~ Z ocp - ) - => QuadratureRoots - -> OcpPhase' ocp - -> OcpPhaseWithCov ocp sx sz sw sr sh shr sc - -> J (CollTraj x z u p n deg) (Vector Double) - -> IO (CollCovProblem ocp n deg sx sw sh shr sc) -makeCollCovProblem roots ocp ocpCov guess = do - let -- the collocation points - taus :: Vec deg Double - taus = mkTaus roots - - computeSensitivities <- mkComputeSensitivities roots (ocpCovDae ocpCov) - computeCovariances <- mkComputeCovariances continuousToDiscreetNoiseApprox - (computeSensitivities) (ocpCovSq ocpCov) - - sbcFun <- toSXFun "sbc" $ \(x0:*:x1) -> ocpCovSbc ocpCov x0 x1 - shFun <- toSXFun "sh" $ \(x0:*:x1) -> ocpCovSh ocpCov (sxSplitJV x0) x1 - mayerFun <- toSXFun "cov mayer" $ \(x0:*:x1:*:x2:*:x3:*:x4) -> - sxCatJV $ Id $ ocpCovMayer ocpCov (unId (sxSplitJV x0)) (sxSplitJV x1) (sxSplitJV x2) x3 x4 - lagrangeFun <- toSXFun "cov lagrange" $ \(x0:*:x1:*:x2:*:x3) -> - sxCatJV $ Id $ ocpCovLagrange ocpCov (unId (sxSplitJV x0)) (sxSplitJV x1) x2 (unId (sxSplitJV x3)) - - cp0 <- makeCollProblem roots ocp guess - - robustify <- mkRobustifyFunction (ocpCovProjection ocpCov) (ocpCovRobustifyPathC ocpCov) - - let nlp0 = cpNlp cp0 - gammas' = ocpCovGammas ocpCov :: shr Double - - gammas :: J (JV shr) MX - gammas = catJV' (fmap realToFrac gammas') - - rpathCUb :: shr Bounds - rpathCUb = fill (Nothing, Just 0) - - robustPathCUb :: J (JV shr) (Vector Bounds) - robustPathCUb = catJV rpathCUb - - -- the NLP - fg :: J (CollTrajCov sx ocp n deg) MX - -> J JNone MX - -> (J (JV Id) MX, J (CollOcpCovConstraints ocp n deg sh shr sc) MX) - fg = getFgCov taus - computeCovariances - gammas - (robustify :: (J (JV shr) MX -> J (JV p) MX -> J (JV x) MX -> J (Cov (JV sx)) MX -> J (JV shr) MX)) - (sbcFun :: SXFun (J (Cov (JV sx)) :*: J (Cov (JV sx))) (J sc)) - (shFun :: SXFun (J (JV x) :*: J (Cov (JV sx))) (J sh)) - (lagrangeFun :: SXFun (J (JV Id) :*: J (JV x) :*: J (Cov (JV sx)) :*: J (JV Id)) (J (JV Id))) - (mayerFun :: SXFun (J (JV Id) :*: (J (JV x) :*: (J (JV x) :*: (J (Cov (JV sx)) :*: J (Cov (JV sx)))))) (J (JV Id))) - (nlpFG nlp0) - - computeCovariancesFun' <- toMXFun "compute covariances" computeCovariances - -- callbacks - let getPlotPoints :: J (CollTrajCov sx ocp n deg) (Vector Double) -> IO (DynPlotPoints Double) - getPlotPoints collTrajCov = do - let CollTrajCov _ collTraj = split collTrajCov - cpPlotPoints cp0 collTraj - - getOutputs :: J (CollTrajCov sx ocp n deg) (Vector Double) - -> IO ( Vec n (Vec deg (o Double, x Double, h Double), x Double) - , Vec n (J (Cov (JV sx)) (Vector Double)) - , J (Cov (JV sx)) (Vector Double) - ) - getOutputs collTrajCov = do - let CollTrajCov _ collTraj = split collTrajCov - outputs <- (cpOutputs cp0) collTraj - covTraj <- fmap split $ eval computeCovariancesFun' (v2d collTrajCov) - let covs' = ctAllButLast covTraj - pF = ctLast covTraj - let covs = unJVec (split covs') :: Vec n (J (Cov (JV sx)) DMatrix) - return (outputs, fmap d2v covs, d2v pF) - - nlp = - Nlp - { nlpFG = fg - , nlpBX = cat $ CollTrajCov (ocpCovS0bnd ocpCov) (nlpBX nlp0) - , nlpBG = cat $ CollOcpCovConstraints - { cocNormal = nlpBG nlp0 - , cocCovPathC = jreplicate (ocpCovShBnds ocpCov) - , cocCovRobustPathC = jreplicate robustPathCUb - , cocSbc = ocpCovSbcBnds ocpCov - } - , nlpX0 = cat $ CollTrajCov (jfill 0) (nlpX0 nlp0) - , nlpP = cat JNone - , nlpLamX0 = Nothing - , nlpLamG0 = Nothing - , nlpScaleF = ocpObjScale ocp - , nlpScaleX = Just $ cat $ - CollTrajCov (fromMaybe (jfill 1) (ocpCovSScale ocpCov)) $ - cat $ fillCollTraj - (fromMaybe (fill 1) (ocpXScale ocp)) - (fromMaybe (fill 1) (ocpZScale ocp)) - (fromMaybe (fill 1) (ocpUScale ocp)) - (fromMaybe (fill 1) (ocpPScale ocp)) - (fromMaybe 1 (ocpTScale ocp)) - - , nlpScaleG = Just $ cat $ CollOcpCovConstraints - { cocNormal = cat $ fillCollConstraints - (fromMaybe (fill 1) (ocpXScale ocp)) - (fromMaybe (fill 1) (ocpResidualScale ocp)) - (fromMaybe (fill 1) (ocpBcScale ocp)) - (fromMaybe (fill 1) (ocpPathCScale ocp)) - , cocCovPathC = jreplicate (fromMaybe (jfill 1) (ocpCovPathCScale ocpCov)) - , cocCovRobustPathC = jreplicate $ - fromMaybe (jfill 1) $ - fmap catJV (ocpCovRobustPathCScale ocpCov) - , cocSbc = fromMaybe (jfill 1) (ocpCovSbcScale ocpCov) - } - } - computeSensitivitiesFun' <- toMXFun "compute sensitivities" computeSensitivities - return $ CollCovProblem { ccpNlp = nlp - , ccpPlotPoints = getPlotPoints - , ccpOutputs = getOutputs - , ccpSensitivities = computeSensitivitiesFun' - , ccpCovariances = computeCovariancesFun' - , ccpRoots = roots - } - -getFg :: - forall x z u p r o c h q n deg . - ( Dim deg, Dim n - , Vectorize x, Vectorize z, Vectorize u, Vectorize p - , Vectorize r, Vectorize o, Vectorize c, Vectorize h, Vectorize q - ) - -- taus - => Vec deg Double - -- bcFun - -> SXFun (J (JV x) :*: J (JV x) :*: J (JV q) :*: J (JV p) :*: J (JV Id)) (J (JV c)) - -- mayerFun - -> SXFun - (J (JV Id) :*: J (JV x) :*: J (JV x) :*: J (JV q) :*: J (JV p)) (J (JV Id)) - -- lagQuadFun - -> ((J (JV p) :*: J (JVec deg (CollPoint (JV x) (JV z) (JV u))) :*: J (JVec deg (JV o)) :*: J (JV Id) :*: J (JVec deg (JV Id))) MX -> - (J (JV Id)) MX) - -- quadFun - -> ((J (JV p) :*: J (JVec deg (CollPoint (JV x) (JV z) (JV u))) :*: J (JVec deg (JV o)) :*: J (JV Id) :*: J (JVec deg (JV Id))) MX -> - (J (JV q)) MX) - -- stageFun - -> ((J (JV Id) :*: J (JV p) :*: J (JVec deg (JV Id)) :*: J (JV x) :*: J (JVec deg (JTuple (JV x) (JV z))) :*: J (JVec deg (JV u))) MX -> (J (JVec deg (JV r)) :*: J (JVec deg (JV o)) :*: J (JVec deg (JV h)) :*: J (JV x)) MX) - -- collTraj - -> J (CollTraj x z u p n deg) MX - -- parameter - -> J JNone MX - -- (objective, constraints) - -> (J (JV Id) MX, J (CollOcpConstraints x r c h n deg) MX) -getFg taus bcFun mayerFun lagQuadFun quadFun stageFun collTraj _ = (obj, cat g) - where - -- split up the design vars - CollTraj tf parm stages' xf = split collTraj - stages = unJVec (split stages') :: Vec n (J (CollStage (JV x) (JV z) (JV u) deg) MX) - spstages = fmap split stages :: Vec n (CollStage (JV x) (JV z) (JV u) deg MX) - - spstagesPoints :: Vec n (J (JVec deg (CollPoint (JV x) (JV z) (JV u))) MX) - spstagesPoints = fmap (\(CollStage _ cps) -> cps) spstages - - obj = objLagrange + objMayer - - objMayer = call mayerFun (tf :*: x0 :*: xf :*: finalQuadratures :*: parm) - - objLagrange :: J (JV Id) MX - objLagrange = F.sum $ TV.tvzipWith3 (oneStage lagQuadFun) spstagesPoints outputs times' - - finalQuadratures :: J (JV q) MX - finalQuadratures = F.sum $ TV.tvzipWith3 (oneStage quadFun) spstagesPoints outputs times' - oneStage :: View qOrSomething - => ((J (JV p) :*: J (JVec deg (CollPoint (JV x) (JV z) (JV u))) :*: J (JVec deg (JV o)) - :*: J (JV Id) :*: J (JVec deg (JV Id))) MX - -> J qOrSomething MX) - -> J (JVec deg (CollPoint (JV x) (JV z) (JV u))) MX - -> J (JVec deg (JV o)) MX - -> J (JVec deg (JV Id)) MX - -> J qOrSomething MX - oneStage qfun stagePoints stageOutputs stageTimes = - qfun (parm :*: stagePoints :*: stageOutputs :*: dt :*: stageTimes) - - -- timestep - dt = tf / fromIntegral n - n = reflectDim (Proxy :: Proxy n) - - -- times at each collocation point - times :: Vec n (Vec deg (J (JV Id) MX)) - times = fmap snd $ timesFromTaus 0 (fmap realToFrac taus) dt - - times' :: Vec n (J (JVec deg (JV Id)) MX) - times' = fmap (cat . JVec) times - - -- initial point at each stage - x0s :: Vec n (J (JV x) MX) - x0s = fmap (\(CollStage x0' _) -> x0') spstages - - -- final point at each stage (for matching constraint) - xfs :: Vec n (J (JV x) MX) - xfs = TV.tvshiftl x0s xf - - x0 = (\(CollStage x0' _) -> x0') (TV.tvhead spstages) - g = CollOcpConstraints - { coCollPoints = cat $ JVec dcs - , coContinuity = cat $ JVec integratorMatchingConstraints - , coPathC = cat $ JVec hs - , coBc = call bcFun (x0 :*: xf :*: finalQuadratures :*: parm :*: tf) - } - - integratorMatchingConstraints :: Vec n (J (JV x) MX) -- THIS SHOULD BE A NONLINEAR FUNCTION - integratorMatchingConstraints = vzipWith (-) interpolatedXs xfs - - dcs :: Vec n (J (JVec deg (JV r)) MX) - outputs :: Vec n (J (JVec deg (JV o)) MX) - hs :: Vec n (J (JVec deg (JV h)) MX) - interpolatedXs :: Vec n (J (JV x) MX) - (dcs, outputs, hs, interpolatedXs) = TV.tvunzip4 $ fmap fff $ TV.tvzip spstages times' - fff :: (CollStage (JV x) (JV z) (JV u) deg MX, J (JVec deg (JV Id)) MX) -> - (J (JVec deg (JV r)) MX, J (JVec deg (JV o)) MX, J (JVec deg (JV h)) MX, J (JV x) MX) - fff (CollStage x0' xzus, stageTimes) = (dc, output, stageHs, interpolatedX') - where - dc :*: output :*: stageHs :*: interpolatedX' = - stageFun (dt :*: parm :*: stageTimes :*: x0' :*: xzs :*: us) - - xzs = cat (JVec xzs') :: J (JVec deg (JTuple (JV x) (JV z))) MX - us = cat (JVec us') :: J (JVec deg (JV u)) MX - (xzs', us') = TV.tvunzip $ fmap toTuple $ unJVec (split xzus) - toTuple xzu = (cat (JTuple x z), u) - where - CollPoint x z u = split xzu - - -getFgCov :: - forall ocp x z u p r c h sx sh shr sc n deg . - ( Dim deg, Dim n, Vectorize x, Vectorize z, Vectorize u, Vectorize p - , Vectorize h, Vectorize c, Vectorize r - , Vectorize sx, View sc, View sh, Vectorize shr - , X ocp ~ x - , Z ocp ~ z - , U ocp ~ u - , P ocp ~ p - , R ocp ~ r - , C ocp ~ c - , H ocp ~ h - ) - -- taus - => Vec deg Double - -> (J (CollTrajCov sx ocp n deg) MX -> J (CovTraj sx n) MX) - -- gammas - -> J (JV shr) MX - -- robustify - -> (J (JV shr) MX -> J (JV p) MX -> J (JV x) MX -> J (Cov (JV sx)) MX -> J (JV shr) MX) - -- sbcFun - -> SXFun (J (Cov (JV sx)) :*: J (Cov (JV sx))) (J sc) - -- shFun - -> SXFun (J (JV x) :*: J (Cov (JV sx))) (J sh) - -- lagrangeFun - -> SXFun - (J (JV Id) :*: J (JV x) :*: J (Cov (JV sx)) :*: J (JV Id)) (J (JV Id)) - -- mayerFun - -> SXFun - (J (JV Id) :*: J (JV x) :*: J (JV x) :*: J (Cov (JV sx)) :*: J (Cov (JV sx))) (J (JV Id)) - -> (J (CollTraj' ocp n deg) MX -> J JNone MX -> (J (JV Id) MX, J (CollOcpConstraints' ocp n deg) MX) - ) - -> J (CollTrajCov sx ocp n deg) MX - -> J JNone MX - -> (J (JV Id) MX, J (CollOcpCovConstraints ocp n deg sh shr sc) MX) -getFgCov - taus computeCovariances - gammas robustify sbcFun shFun lagrangeFun mayerFun - normalFG collTrajCov nlpParams = - (obj0 + objectiveLagrangeCov + objectiveMayerCov, cat g) - where - CollTrajCov p0 collTraj = split collTrajCov - (obj0, g0) = normalFG collTraj nlpParams - - g = CollOcpCovConstraints - { cocNormal = g0 - , cocCovPathC = cat (JVec covPathConstraints) - , cocCovRobustPathC = cat (JVec robustifiedPathC) - , cocSbc = call sbcFun (p0 :*: pF) - } - -- split up the design vars - CollTraj tf parm stages' xf = split collTraj - stages = unJVec (split stages') :: Vec n (J (CollStage (JV x) (JV z) (JV u) deg) MX) - spstages = fmap split stages :: Vec n (CollStage (JV x) (JV z) (JV u) deg MX) - - objectiveMayerCov = call mayerFun (tf :*: x0 :*: xf :*: p0 :*: pF) - - -- timestep - dt = tf / fromIntegral n - n = reflectDim (Proxy :: Proxy n) - - -- times at each collocation point - t0s :: Vec n (J (JV Id) MX) - (t0s, _) = TV.tvunzip $ timesFromTaus 0 (fmap realToFrac taus) dt - - -- initial point at each stage - x0s :: Vec n (J (JV x) MX) - x0s = fmap (\(CollStage x0' _) -> x0') spstages - - x0 = (\(CollStage x0' _) -> x0') (TV.tvhead spstages) - --- sensitivities = call computeSensitivities collTraj - - covs :: Vec n (J (Cov (JV sx)) MX) - covs = unJVec (split covs') - - covs' :: J (JVec n (Cov (JV sx))) MX -- all but last covariance - pF :: J (Cov (JV sx)) MX -- last covariances - CovTraj covs' pF = split (computeCovariances collTrajCov) - - -- lagrange term - objectiveLagrangeCov = (lagrangeF + lagrange0s) / fromIntegral n - where - lagrangeF = call lagrangeFun (tf :*: xf :*: pF :*: tf) - lagrange0s = - sum $ F.toList $ - TV.tvzipWith3 (\tk xk pk -> call lagrangeFun (tk :*: xk :*: pk :*: tf)) t0s x0s covs - - covPathConstraints :: Vec n (J sh MX) - covPathConstraints = TV.tvzipWith (\xk pk -> call shFun (xk:*:pk)) x0s covs - - robustifiedPathC :: Vec n (J (JV shr) MX) - robustifiedPathC = TV.tvzipWith (robustify gammas parm) x0s covs - - -getBg :: forall x z u p r o c h q n deg . - ( Dim n, Dim deg - , Vectorize x, Vectorize r, Vectorize c, Vectorize h - ) - => OcpPhase x z u p r o c h q - -> CollOcpConstraints x r c h n deg (Vector Bounds) -getBg ocp = - CollOcpConstraints - { coCollPoints = jreplicate (jfill (Just 0, Just 0)) -- dae residual constraint - , coContinuity = jreplicate (jfill (Just 0, Just 0)) -- continuity constraint - , coPathC = jreplicate (jreplicate hbnds) - , coBc = catJV (ocpBcBnds ocp) - } - where - hbnds :: J (JV h) (Vector Bounds) - hbnds = catJV (ocpPathCBnds ocp) - -evaluateQuadraturesFunction :: - forall x z u p o q deg . - (Dim deg, View x, View z, View u, View o, View p, View q) - => SXFun (J x :*: J z :*: J u :*: J p :*: J o :*: J (JV Id) :*: J (JV Id)) (J q) - -> (J q MX -> Vec deg (J q MX) -> J q MX) - -> Vec (TV.Succ deg) (Vec (TV.Succ deg) Double) - -> Int - -> (J p :*: J (JVec deg (CollPoint x z u)) :*: J (JVec deg o) :*: J (JV Id) :*: J (JVec deg (JV Id))) MX - -> J q MX -evaluateQuadraturesFunction f interpolate' cijs' n (p :*: stage' :*: outputs' :*: dt :*: stageTimes') = - M.uncol $ M.ms (M.col qnext) dt - where - tf = dt * fromIntegral n - - stage :: Vec deg (CollPoint x z u MX) - stage = fmap split $ unJVec $ split stage' - - outputs :: Vec deg (J o MX) - outputs = unJVec (split outputs') - - stageTimes :: Vec deg (J (JV Id) MX) - stageTimes = unJVec (split stageTimes') - - qdots :: Vec deg (J q MX) - qdots = TV.tvzipWith3 (\(CollPoint x z u) o t -> call f (x:*:z:*:u:*:p:*:o:*:t:*:tf)) stage outputs stageTimes - - qnext :: J q MX - qnext = interpolate' (0 :: J q MX) qs - - qs :: Vec deg (J q MX) - qs = cijInvFr !* qdots - - cijs :: Vec deg (Vec deg Double) - cijs = TV.tvtail $ fmap TV.tvtail cijs' - - cijMat :: Mat.Matrix Double - cijMat = Mat.fromLists $ F.toList $ fmap F.toList cijs - - cijInv' :: Mat.Matrix Double - cijInv' = LA.inv cijMat - - cijInv :: Vec deg (Vec deg Double) - cijInv = TV.mkVec' (map TV.mkVec' (Mat.toLists cijInv')) - - cijInvFr :: Vec deg (Vec deg (J q MX)) - cijInvFr = fmap (fmap realToFrac) cijInv - - --- todo: merging this with evaluateQuadraturesFunction would reduce duplication, --- but could be inefficient -genericQuadraturesFunction :: - forall deg - . Dim deg - => (J (JV Id) MX -> Vec deg (J (JV Id) MX) -> J (JV Id) MX) - -> Vec (TV.Succ deg) (Vec (TV.Succ deg) Double) - -> Int - -> (J (JVec deg (JV Id)) :*: J (JV Id)) MX - -> J (JV Id) MX -genericQuadraturesFunction interpolate' cijs' n (qdots' :*: tf) = - dt * qnext - where - dt = tf / fromIntegral n - - qdots :: Vec deg (J (JV Id) MX) - qdots = unJVec $ split qdots' - - qnext :: J (JV Id) MX - qnext = interpolate' 0 qs - - qs = cijInvFr !* qdots - - cijs :: Vec deg (Vec deg Double) - cijs = TV.tvtail $ fmap TV.tvtail cijs' - - cijMat :: Mat.Matrix Double - cijMat = Mat.fromLists $ F.toList $ fmap F.toList cijs - - cijInv' :: Mat.Matrix Double - cijInv' = LA.inv cijMat - - cijInv :: Vec deg (Vec deg Double) - cijInv = TV.mkVec' (map TV.mkVec' (Mat.toLists cijInv')) - - cijInvFr :: Vec deg (Vec deg (J (JV Id) MX)) - cijInvFr = fmap (fmap realToFrac) cijInv - - --- todo: code duplication -dot :: forall x deg a b. (Fractional (J x a), Real b, Dim deg) => Vec deg b -> Vec deg (J x a) -> J x a -dot cks xs = F.sum $ TV.unVec elemwise - where - elemwise :: Vec deg (J x a) - elemwise = TV.tvzipWith smul cks xs - - smul :: b -> J x a -> J x a - smul x y = realToFrac x * y - - --- todo: code duplication -interpolateXDots' :: (Real b, Fractional (J x a), Dim deg) => Vec deg (Vec deg b) -> Vec deg (J x a) -> Vec deg (J x a) -interpolateXDots' cjks xs = fmap (`dot` xs) cjks - -interpolateXDots :: - (Real b, Dim deg, Fractional (J x a)) => - Vec (TV.Succ deg) (Vec (TV.Succ deg) b) - -> Vec (TV.Succ deg) (J x a) - -> Vec deg (J x a) -interpolateXDots cjks xs = TV.tvtail $ interpolateXDots' cjks xs - - --- path constraints -pathConFunction :: - forall x z u p o h a . (View x, View z, View u, View o, View h, Viewable a) - => (J x a -> J z a -> J u a -> J p a -> J o a -> J (JV Id) a -> J h a) - -> (J (JV Id) :*: J p :*: J o :*: J (CollPoint x z u)) a - -> J h a -pathConFunction pathC (t :*: parm :*: o :*: collPoint) = - pathC x z u parm o t - where - CollPoint x z u = split collPoint - --- return dynamics constraints, outputs, and interpolated state -dynStageConstraints :: - forall x z u p r o deg . (Dim deg, View x, View z, View u, View p, View r, View o) - => (J x MX -> Vec deg (J x MX) -> J x MX) - -> Vec (TV.Succ deg) (Vec (TV.Succ deg) Double) - -> SXFun (J (JV Id) :*: J p :*: J x :*: J (CollPoint x z u)) - (J r :*: J o) - -> (J x :*: J (JVec deg (JTuple x z)) :*: J (JVec deg u) :*: J (JV Id) :*: J p :*: J (JVec deg (JV Id))) MX - -> (J (JVec deg r) :*: J x :*: J (JVec deg o)) MX -dynStageConstraints interpolate' cijs dynFun (x0 :*: xzs' :*: us' :*: h :*: p :*: stageTimes') = - cat (JVec dynConstrs) :*: xnext :*: cat (JVec outputs) - where - xzs = fmap split (unJVec (split xzs')) :: Vec deg (JTuple x z MX) - us = unJVec (split us') :: Vec deg (J u MX) - - -- interpolated final state - xnext :: J x MX - xnext = interpolate' x0 xs - - stageTimes = unJVec $ split stageTimes' - - -- dae constraints (dynamics) - dynConstrs :: Vec deg (J r MX) - outputs :: Vec deg (J o MX) - (dynConstrs, outputs) = TV.tvunzip $ TV.tvzipWith4 applyDae xdots xzs us stageTimes - - applyDae :: J x MX -> JTuple x z MX -> J u MX -> J (JV Id) MX -> (J r MX, J o MX) - applyDae x' (JTuple x z) u t = (r, o) - where - r :*: o = call dynFun (t :*: p :*: x' :*: collPoint) - collPoint = cat (CollPoint x z u) - - -- state derivatives, maybe these could be useful as outputs - xdots :: Vec deg (J x MX) - xdots = fmap (`M.vs` (1/h)) $ interpolateXDots cijs (x0 TV.<| xs) - - xs :: Vec deg (J x MX) - xs = fmap (\(JTuple x _) -> x) xzs - - -data ErrorIn0 x z u p deg a = - ErrorIn0 (J x a) (J (JVec deg (CollPoint x z u)) a) (J (JV Id) a) (J p a) (J (JVec deg (JV Id)) a) - deriving Generic -data ErrorInD sx sw sz deg a = - ErrorInD (J sx a) (J sw a) (J (JVec deg (JTuple sx sz)) a) - deriving Generic -data ErrorOut sr sx deg a = - ErrorOut (J (JVec deg sr) a) (J sx a) - deriving Generic - -instance (View x, View z, View u, View p, Dim deg) => Scheme (ErrorIn0 x z u p deg) -instance (View sx, View sw, View sz, Dim deg) => View (ErrorInD sx sw sz deg) -instance (View sr, View sx, Dim deg) => View (ErrorOut sr sx deg) - - - --- outputs -outputFunction :: - forall x z u p r o deg . (Dim deg, View x, View z, View u, View p, View r, View o) - => (J x MX -> Vec deg (J x MX) -> J x MX) - -> Vec (TV.Succ deg) (Vec (TV.Succ deg) Double) -> Vec deg Double - -> SXFun (J (JV Id) :*: J p :*: J x :*: J (CollPoint x z u)) - (J r :*: J o) - -> (J (CollStage x z u deg) :*: J p :*: J (JV Id) :*: J (JV Id)) MX - -> (J (JVec deg r) :*: J (JVec deg x) :*: J (JVec deg o) :*: J x) MX -outputFunction callInterpolate cijs taus dynFun (collStage :*: p :*: h :*: k) = - cat (JVec dynConstrs) :*: cat (JVec xdots) :*: cat (JVec outputs) :*: xnext - where - xzus = unJVec (split xzus') :: Vec deg (J (CollPoint x z u) MX) - CollStage x0 xzus' = split collStage - -- times at each collocation point - stageTimes :: Vec deg (J (JV Id) MX) - stageTimes = fmap (\tau -> t0 + realToFrac tau * h) taus - t0 = k*h - - xnext = callInterpolate x0 xs - - -- dae constraints (dynamics) - dynConstrs :: Vec deg (J r MX) - outputs :: Vec deg (J o MX) - (dynConstrs, outputs) = TV.tvunzip $ TV.tvzipWith3 applyDae xdots xzus stageTimes - - applyDae :: J x MX -> J (CollPoint x z u) MX -> J (JV Id) MX -> (J r MX, J o MX) - applyDae x' xzu t = (r, o) - where - r :*: o = call dynFun (t :*: p :*: x' :*: xzu) - - -- state derivatives, maybe these could be useful as outputs - xdots :: Vec deg (J x MX) - xdots = fmap (`M.vs` (1/h)) $ interpolateXDots cijs (x0 TV.<| xs) - - xs :: Vec deg (J x MX) - xs = fmap ((\(CollPoint x _ _) -> x) . split) xzus - - - --- return path constraints at each collocation point -pathStageConstraints :: - forall x z u p o h deg . (Dim deg, View x, View z, View u, View p, View o, View h) - => SXFun (J (JV Id) :*: J p :*: J o :*: J (CollPoint x z u)) - (J h) - -> (J p :*: J (JVec deg (JV Id)) :*: J (JVec deg o) :*: J (JVec deg (CollPoint x z u))) MX - -> J (JVec deg h) MX -pathStageConstraints pathCFun - (p :*: stageTimes' :*: outputs :*: collPoints) = - cat (JVec hs) - where - stageTimes = unJVec $ split stageTimes' - cps = fmap split (unJVec (split collPoints)) :: Vec deg (CollPoint x z u MX) - - -- path constraints - hs :: Vec deg (J h MX) - hs = TV.tvzipWith3 applyH cps stageTimes (unJVec (split outputs)) - - applyH :: CollPoint x z u MX -> J (JV Id) MX -> J o MX -> J h MX - applyH (CollPoint x z u) t o = pathc' - where - pathc' = call pathCFun (t :*: p :*: o :*: collPoint) - collPoint = cat (CollPoint x z u) - - -stageFunction :: - forall x z u p o r h deg . (Dim deg, View x, View z, View u, View p, View r, View o, View h) - => MXFun (J p :*: J (JVec deg (JV Id)) :*: J (JVec deg o) :*: J (JVec deg (CollPoint x z u))) - (J (JVec deg h)) - -> ((J x :*: J (JVec deg (JTuple x z)) :*: J (JVec deg u) :*: J (JV Id) :*: J p :*: J (JVec deg (JV Id))) MX - -> (J (JVec deg r) :*: J x :*: J (JVec deg o)) MX) - -> (J (JV Id) :*: J p :*: J (JVec deg (JV Id)) :*: J x :*: J (JVec deg (JTuple x z)) :*: J (JVec deg u)) MX - -> (J (JVec deg r) :*: J (JVec deg o) :*: J (JVec deg h) :*: J x) MX -stageFunction pathConStageFun dynStageCon - (dt :*: parm :*: stageTimes :*: x0' :*: xzs' :*: us) = - dynConstrs :*: outputs :*: hs :*: interpolatedX - where - collPoints = cat $ JVec $ TV.tvzipWith catXzu (unJVec (split xzs')) (unJVec (split us)) - - catXzu :: J (JTuple x z) MX -> J u MX -> J (CollPoint x z u) MX - catXzu xz u = cat $ CollPoint x z u - where - JTuple x z = split xz - - dynConstrs :: J (JVec deg r) MX - outputs :: J (JVec deg o) MX - interpolatedX :: J x MX - (dynConstrs :*: interpolatedX :*: outputs) = - dynStageCon (x0' :*: xzs' :*: us :*: dt :*: parm :*: stageTimes) - - hs :: J (JVec deg h) MX - hs = call pathConStageFun (parm :*: stageTimes :*: outputs :*: collPoints) - - --- | make an initial guess -makeGuess :: - forall x z u p deg n . - ( Dim n, Dim deg - , Vectorize x, Vectorize z, Vectorize u, Vectorize p - ) - => QuadratureRoots - -> Double -> (Double -> x Double) -> (Double -> z Double) -> (Double -> u Double) - -> p Double - -> CollTraj x z u p n deg (Vector Double) -makeGuess quadratureRoots tf guessX guessZ guessU parm = - CollTraj (jfill tf) (catJV parm) guesses (catJV (guessX tf)) - where - -- timestep - dt = tf / fromIntegral n - n = vlength (Proxy :: Proxy (Vec n)) - - -- initial time at each collocation stage - t0s :: Vec n Double - t0s = TV.mkVec' $ take n [dt * fromIntegral k | k <- [(0::Int)..]] - - -- times at each collocation point - times :: Vec n (Double, Vec deg Double) - times = fmap (\t0 -> (t0, fmap (\tau -> t0 + tau*dt) taus)) t0s - - mkGuess' :: (Double, Vec deg Double) -> CollStage (JV x) (JV z) (JV u) deg (Vector Double) - mkGuess' (t,ts) = - CollStage (catJV (guessX t)) $ - cat $ JVec $ fmap (\t' -> cat (CollPoint (catJV (guessX t')) (catJV (guessZ t')) (catJV (guessU t')))) ts - - guesses :: J (JVec n (CollStage (JV x) (JV z) (JV u) deg)) (Vector Double) - guesses = cat $ JVec $ fmap (cat . mkGuess') times - - -- the collocation points - taus :: Vec deg Double - taus = mkTaus quadratureRoots - - --- | make an initial guess -makeGuessSim :: - forall x z u p deg n . - ( Dim n, Dim deg - , Vectorize x, Vectorize z, Vectorize u, Vectorize p - ) - => QuadratureRoots - -> Double - -> x Double - -> (x Double -> u Double -> x Double) - -> (x Double -> Double -> u Double) - -> p Double - -> CollTraj x z u p n deg (Vector Double) -makeGuessSim quadratureRoots tf x00 ode guessU p = - CollTraj (jfill tf) (catJV p) (cat (JVec stages)) (catJV xf) - where - -- timestep - dt = tf / fromIntegral n - n = vlength (Proxy :: Proxy (Vec n)) - - -- initial time at each collocation stage - t0s :: Vec n Double - t0s = TV.mkVec' $ take n [dt * fromIntegral k | k <- [(0::Int)..]] - - xf :: x Double - stages :: Vec n (J (CollStage (JV x) (JV z) (JV u) deg) (Vector Double)) - (xf, stages) = T.mapAccumL stageGuess x00 t0s - - stageGuess :: x Double -> Double - -> (x Double, J (CollStage (JV x) (JV z) (JV u) deg) (Vector Double)) - stageGuess x0 t0 = (integrate 1, cat (CollStage (catJV x0) points)) - where - points = cat $ JVec $ fmap (toCollPoint . integrate) taus - u = guessU x0 t0 - f x = ode x u - toCollPoint x = cat $ CollPoint (catJV x) (catJV (fill 0 :: z Double)) (catJV u) - integrate localTau = rk4 f (localTau * dt) x0 - - -- the collocation points - taus :: Vec deg Double - taus = mkTaus quadratureRoots - - rk4 :: (x Double -> x Double) -> Double -> x Double -> x Double - rk4 f h x0 = x0 ^+^ ((k1 ^+^ (2 *^ k2) ^+^ (2 *^ k3) ^+^ k4) ^/ 6) - where - k1 = (f x0) ^* h - k2 = (f (x0 ^+^ (k1^/2))) ^* h - k3 = (f (x0 ^+^ (k2^/2))) ^* h - k4 = (f (x0 ^+^ k3)) ^* h - - (^+^) :: x Double -> x Double -> x Double - y0 ^+^ y1 = devectorize $ V.zipWith (+) (vectorize y0) (vectorize y1) - - (*^) :: Double -> x Double -> x Double - y0 *^ y1 = devectorize $ V.map (y0 *) (vectorize y1) - - (^*) :: x Double -> Double -> x Double - y0 ^* y1 = devectorize $ V.map (* y1) (vectorize y0) - - (^/) :: x Double -> Double -> x Double - y0 ^/ y1 = devectorize $ V.map (/ y1) (vectorize y0) +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE PolyKinds #-} + +module Dyno.DirectCollocation.Formulate + ( CollProblem(..) + , StageOutputs(..) + , makeCollProblem + , mkTaus + , makeGuess + , makeGuessSim + , ocpPhaseBx + , ocpPhaseBg + ) where + +import GHC.Generics ( Generic, Generic1 ) + +import Control.Applicative +import Control.Monad.State ( StateT(..), runStateT ) +import Data.Maybe ( fromMaybe ) +import Data.Proxy ( Proxy(..) ) +import Data.Vector ( Vector ) +import qualified Data.Foldable as F +import qualified Data.Traversable as T +import qualified Numeric.LinearAlgebra as Mat +import Linear hiding ( dot ) +import Prelude -- BBP workaround + +import Casadi.DMatrix ( DMatrix ) +import Casadi.MX ( MX ) +import Casadi.SX ( SX ) + +import Dyno.Integrate ( InitialTime(..), TimeStep(..), rk45 ) +import Dyno.View.View ( View(..), J, jfill, JTuple(..), v2d, d2v ) +import qualified Dyno.View.M as M +import Dyno.View.JV ( JV, splitJV, catJV, splitJV', catJV' ) +import Dyno.View.HList ( (:*:)(..) ) +import Dyno.View.Fun +import Dyno.View.JVec( JVec(..), jreplicate ) +import Dyno.View.Scheme ( Scheme ) +import Dyno.Vectorize ( Vectorize(..), Id(..), fill, vlength, vzipWith ) +import Dyno.TypeVecs ( Vec, Dim, reflectDim ) +import qualified Dyno.TypeVecs as TV +import Dyno.LagrangePolynomials ( lagrangeDerivCoeffs ) +import Dyno.Nlp ( Nlp(..), Bounds ) +import Dyno.Ocp + +import Dyno.DirectCollocation.Types +import Dyno.DirectCollocation.Dynamic ( MetaProxy(..), DynPlotPoints, dynPlotPoints ) +import Dyno.DirectCollocation.Quadratures ( QuadratureRoots(..), mkTaus, interpolate, timesFromTaus ) + +data CollProblem x z u p r o c h q qo po fp n deg = + CollProblem + { cpNlp :: Nlp (CollTraj x z u p n deg) + (JV fp) + (CollOcpConstraints x r c h n deg) MX + , cpOcp :: OcpPhase x z u p r o c h q qo po fp + , cpPlotPoints :: J (CollTraj x z u p n deg) (Vector Double) + -> J (JV fp) (Vector Double) + -> IO (DynPlotPoints Double) + , cpHellaOutputs :: J (CollTraj x z u p n deg) (Vector Double) + -> J (JV fp) (Vector Double) + -> IO ( DynPlotPoints Double + , Vec n (StageOutputs x o h q qo po deg Double) + , Quadratures q qo Double + ) + , cpConstraints :: J (CollTraj x z u p n deg) (Vector Double) + -> J (JV fp) (Vector Double) + -> IO (J (CollOcpConstraints x r c h n deg) (Vector Double)) + , cpOutputs :: J (CollTraj x z u p n deg) (Vector Double) + -> J (JV fp) (Vector Double) + -> IO (Vec n (StageOutputs x o h q qo po deg Double)) + , cpTaus :: Vec deg Double + , cpRoots :: QuadratureRoots + , cpEvalQuadratures :: Vec n (Vec deg Double) -> Double -> IO Double + , cpMetaProxy :: MetaProxy x z u p o q qo po h + } + + +data QuadraturePlottingIn x z u p o q qo fp a = + -- x0 xF x z u p fp o q qo t T + QuadraturePlottingIn (J x a) (J x a) (J x a) (J z a) (J u a) (J p a) (J o a) (J q a) (J qo a) (J fp a) + (J (JV Id) a) (J (JV Id) a) + deriving (Generic, Generic1) + +data QuadratureIn x z u p fp a = + -- x' x z u p fp t T + QuadratureIn (J x a) (J x a) (J z a) (J u a) (J p a) (J fp a) + (J (JV Id) a) (J (JV Id) a) + deriving (Generic, Generic1) + +data QuadratureStageIn x z u p fp deg a = + -- xzus p fp ts h + QuadratureStageIn (J (CollStage x z u deg) a) (J p a) (J fp a) (J (JVec deg (JV Id)) a) (J (JV Id) a) + deriving (Generic, Generic1) + +data QuadratureStageOut q deg a = + -- qdots qs qNext + QuadratureStageOut (J (JVec deg q) a) (J (JVec deg q) a) (J q a) + deriving (Generic, Generic1) + +data PathCIn x z u p fp a = + -- x' x z u p t + PathCIn (J x a) (J x a) (J z a) (J u a) (J p a) (J fp a) (J (JV Id) a) + deriving (Generic, Generic1) + +data PathCStageIn x z u p fp deg a = + -- xzus p fp ts h + PathCStageIn (J (CollStage x z u deg) a) (J p a) (J fp a) (J (JVec deg (JV Id)) a) (J (JV Id) a) + deriving (Generic, Generic1) + +data DaeIn x z u p fp a = + -- t p fp x' (CollPoint x z u) + DaeIn (J (JV Id) a) (J p a) (J fp a) (J x a) (J (CollPoint x z u) a) + deriving (Generic, Generic1) + +data DaeOut r o a = + -- r o + DaeOut (J r a) (J o a) + deriving (Generic, Generic1) + +instance (View x, View z, View u, View p, View o, View q, View qo, View fp) + => Scheme (QuadraturePlottingIn x z u p o q qo fp) +instance (View x, View z, View u, View p, View fp) => Scheme (QuadratureIn x z u p fp) +instance (View x, View z, View u, View p, View fp, Dim deg) => Scheme (QuadratureStageIn x z u p fp deg) +instance (View q, Dim deg) => Scheme (QuadratureStageOut q deg) +instance (View x, View z, View u, View p, View fp) => Scheme (PathCIn x z u p fp) +instance (View x, View z, View u, View p, View fp, Dim deg) => Scheme (PathCStageIn x z u p fp deg) +instance (View x, View z, View u, View p, View fp) => Scheme (DaeIn x z u p fp) +instance (View r, View o) => Scheme (DaeOut r o) + + +--toQuadratureOcp :: (Vectorize x, Vectorize q, Vectorize c, Vectorize r) +-- => q Double +-- -> q Double +-- -> q Double +-- -> OcpPhase x z u p r o c h q qo po fp +-- -> OcpPhase (Tuple x q) z u p (Tuple r q) o (Tuple c q) h None qo po fp +--toQuadratureOcp qscale q0scale qdotScale ocp0 = +-- OcpPhase +-- { ocpMayer = \tf (Tuple x0 _) (Tuple xf qf) None p fp -> ocpMayer ocp0 tf x0 xf qf p fp +-- , ocpLagrange = \(Tuple x _) z u p fp o t -> ocpLagrange ocp0 x z u p fp o t +-- , ocpQuadratures = \_ _ _ _ _ _ _ _ -> None +-- , ocpQuadratureOutputs = \(Tuple x _) -> ocpQuadratureOutputs ocp0 x +-- , ocpDae = \(Tuple x' q') (Tuple x _) z u p fp t -> +-- let (res0, o) = ocpDae ocp0 x' x z u p fp t +-- tf = error "toQuadratureOcp: quadrature derivatives can't use end time" +-- dq = ocpQuadratures ocp0 x z u p fp o t tf +-- in (Tuple res0 (vzipWith (-) q' dq), o) +-- , ocpBc = \(Tuple x0 q0) (Tuple xf qf) None p fp tf -> +-- Tuple (ocpBc ocp0 x0 xf qf p fp tf) q0 +-- , ocpPathC = \(Tuple x _) -> ocpPathC ocp0 x +-- , ocpPlotOutputs = \(Tuple x q) z u p o None -> ocpPlotOutputs ocp0 x z u p o q +-- , ocpObjScale = ocpObjScale ocp0 +-- , ocpTScale = ocpTScale ocp0 +-- , ocpXScale = Just $ Tuple (fromMaybe (fill 1) (ocpXScale ocp0)) qscale +-- , ocpZScale = ocpZScale ocp0 +-- , ocpUScale = ocpUScale ocp0 +-- , ocpPScale = ocpPScale ocp0 +-- , ocpResidualScale = Just $ Tuple (fromMaybe (fill 1) (ocpResidualScale ocp0)) qdotScale +-- , ocpBcScale = Just $ Tuple (fromMaybe (fill 1) (ocpBcScale ocp0)) q0scale +-- , ocpPathCScale = ocpPathCScale ocp0 +-- } + +makeCollProblem :: + forall x z u p r o c h q qo po fp deg n . + ( Dim deg, Dim n + , Vectorize x, Vectorize p, Vectorize u, Vectorize z + , Vectorize r, Vectorize o, Vectorize h, Vectorize c, Vectorize q + , Vectorize po, Vectorize fp, Vectorize qo + ) + => QuadratureRoots + -> OcpPhase x z u p r o c h q qo po fp + -> OcpPhaseInputs x z u p c h fp + -> J (CollTraj x z u p n deg) (Vector Double) + -> IO (CollProblem x z u p r o c h q qo po fp n deg) +makeCollProblem roots ocp ocpInputs guess = do + let -- the collocation points + taus :: Vec deg Double + taus = mkTaus roots + + n = reflectDim (Proxy :: Proxy n) + + -- coefficients for getting xdot by lagrange interpolating polynomials + cijs :: Vec (TV.Succ deg) (Vec (TV.Succ deg) Double) + cijs = lagrangeDerivCoeffs (0 TV.<| taus) + + interpolate' :: View f => (J f :*: J (JVec deg f)) MX -> J f MX + interpolate' (x0 :*: xs) = case roots of + Legendre -> interpolate taus x0 (unJVec (split xs)) + Radau -> TV.tvlast $ unJVec $ split xs + + dynamicsFunction :: DaeIn (JV x) (JV z) (JV u) (JV p) (JV fp) SX -> DaeOut (JV r) (JV o) SX + dynamicsFunction (DaeIn t parm fixedParm x' collPoint) = DaeOut (catJV' r) (catJV' o) + where + CollPoint x z u = split collPoint + (r,o) = ocpDae ocp + (splitJV' x') (splitJV' x) (splitJV' z) (splitJV' u) + (splitJV' parm) (splitJV' fixedParm) (unId (splitJV' t)) + + interpolateFun <- toMXFun "interpolate (JV x)" interpolate' >>= expandMXFun + interpolateQFun <- toMXFun "interpolate (JV q)" interpolate' >>= expandMXFun + interpolateQoFun <- toMXFun "interpolate (JV qo)" interpolate' >>= expandMXFun + interpolateScalarFun <- toMXFun "interpolate (JV Id)" interpolate' >>= expandMXFun + + let callInterpolateScalar :: J (JV Id) MX -> Vec deg (J (JV Id) MX) -> J (JV Id) MX + callInterpolateScalar x0 xs = call interpolateScalarFun (x0 :*: cat (JVec xs)) + + callInterpolate :: J (JV x) MX -> Vec deg (J (JV x) MX) -> J (JV x) MX + callInterpolate x0 xs = call interpolateFun (x0 :*: cat (JVec xs)) + + callInterpolateQ :: J (JV q) MX -> Vec deg (J (JV q) MX) -> J (JV q) MX + callInterpolateQ q0 qs = call interpolateQFun (q0 :*: cat (JVec qs)) + + callInterpolateQo :: J (JV qo) MX -> Vec deg (J (JV qo) MX) -> J (JV qo) MX + callInterpolateQo q0 qs = call interpolateQoFun (q0 :*: cat (JVec qs)) + + let quadFun :: QuadratureIn (JV x) (JV z) (JV u) (JV p) (JV fp) SX -> J (JV q) SX + quadFun (QuadratureIn x' x z u p fp t tf) = quad + where + daeIn = DaeIn t p fp x' (cat (CollPoint x z u)) + DaeOut _ o = dynamicsFunction daeIn + + quad :: J (JV q) SX + quad = catJV' $ ocpQuadratures ocp + (splitJV' x) (splitJV' z) (splitJV' u) (splitJV' p) (splitJV' fp) (splitJV' o) + (unId (splitJV' t)) (unId (splitJV' tf)) + + let quadOutFun :: QuadratureIn (JV x) (JV z) (JV u) (JV p) (JV fp) SX -> J (JV qo) SX + quadOutFun (QuadratureIn x' x z u p fp t tf) = quad + where + daeIn = DaeIn t p fp x' (cat (CollPoint x z u)) + DaeOut _ o = dynamicsFunction daeIn + + quad :: J (JV qo) SX + quad = catJV' $ ocpQuadratureOutputs ocp + (splitJV' x) (splitJV' z) (splitJV' u) (splitJV' p) (splitJV' fp) (splitJV' o) + (unId (splitJV' t)) (unId (splitJV' tf)) + + let lagFun :: QuadratureIn (JV x) (JV z) (JV u) (JV p) (JV fp) SX -> J (JV Id) SX + lagFun (QuadratureIn x' x z u p fp t tf) = lag + where + daeIn = DaeIn t p fp x' (cat (CollPoint x z u)) + DaeOut _ o = dynamicsFunction daeIn + + lag :: J (JV Id) SX + lag = catJV' $ Id $ ocpLagrange ocp + (splitJV' x) (splitJV' z) (splitJV' u) (splitJV' p) (splitJV' fp) (splitJV' o) + (unId (splitJV' t)) (unId (splitJV' tf)) + + let pathCFun :: PathCIn (JV x) (JV z) (JV u) (JV p) (JV fp) SX -> J (JV h) SX + pathCFun (PathCIn x' x z u p fp t) = h + where + daeIn = DaeIn t p fp x' (cat (CollPoint x z u)) + DaeOut _ o = dynamicsFunction daeIn + + h :: J (JV h) SX + h = catJV' $ ocpPathC ocp + (splitJV' x) (splitJV' z) (splitJV' u) (splitJV' p) (splitJV' fp) (splitJV' o) + (unId (splitJV' t)) + + quadFunSX <- toSXFun "quadFun" quadFun + quadOutFunSX <- toSXFun "quadOutFun" quadOutFun + lagFunSX <- toSXFun "lagFun" lagFun + pathCFunSX <- toSXFun "pathCFun" pathCFun + + let quadraturePlottingFun :: + QuadraturePlottingIn (JV x) (JV z) (JV u) (JV p) (JV o) (JV q) (JV qo) (JV fp) SX -> J (JV po) SX + quadraturePlottingFun (QuadraturePlottingIn x0 xF x z u p o q qo fp t tf) = + catJV' $ ocpPlotOutputs ocp (splitJV' x0, splitJV' xF) + (splitJV' x) (splitJV' z) (splitJV' u) (splitJV' p) + (splitJV' o) (splitJV' q) (splitJV' qo) (splitJV' fp) + (unId (splitJV' t)) (unId (splitJV' tf)) + quadPlotFunSX <- toSXFun "quadPlotFun" quadraturePlottingFun + + let -- later we could use the intermediate points as outputs, or in path cosntraints + lagrangeStageFun :: QuadratureStageIn (JV x) (JV z) (JV u) (JV p) (JV fp) deg MX + -> QuadratureStageOut (JV Id) deg MX + lagrangeStageFun qIn = QuadratureStageOut (cat (JVec qdots)) (cat (JVec qs)) qNext + where + (qdots,qs,qNext) = toQuadratureFun n cijs callInterpolateScalar (call lagFunSX) qIn + quadratureStageFun :: QuadratureStageIn (JV x) (JV z) (JV u) (JV p) (JV fp) deg MX + -> QuadratureStageOut (JV q) deg MX + quadratureStageFun qIn = QuadratureStageOut (cat (JVec qdots)) (cat (JVec qs)) qNext + where + (qdots,qs,qNext) = toQuadratureFun n cijs callInterpolateQ (call quadFunSX) qIn + quadratureOutStageFun :: QuadratureStageIn (JV x) (JV z) (JV u) (JV p) (JV fp) deg MX + -> QuadratureStageOut (JV qo) deg MX + quadratureOutStageFun qIn = QuadratureStageOut (cat (JVec qdots)) (cat (JVec qs)) qNext + where + (qdots,qs,qNext) = toQuadratureFun n cijs callInterpolateQo (call quadOutFunSX) qIn + pathCStageFun pcIn = cat (JVec hs) + where + hs = toPathCFun cijs (call pathCFunSX) pcIn + lagrangeStageFunMX <- toMXFun "lagrangeStageFun" $ + (\(QuadratureStageOut _ _ q) -> q) . lagrangeStageFun + quadratureStageFunMX <- toMXFun "quadratureStageFun" $ + (\(QuadratureStageOut _ _ q) -> q) . quadratureStageFun + pathCStageFunMX <- toMXFun "pathCStageFun" pathCStageFun + + + bcFun <- toSXFun "bc" $ \(x0:*:x1:*:x2:*:x3:*:x4:*:x5) -> catJV' $ ocpBc ocp (splitJV' x0) (splitJV' x1) (splitJV' x2) (splitJV' x3) (splitJV' x4) (unId (splitJV' x5)) + mayerFun <- toSXFun "mayer" $ \(x0:*:x1:*:x2:*:x3:*:x4:*:x5) -> + catJV' $ Id $ ocpMayer ocp (unId (splitJV' x0)) (splitJV' x1) (splitJV' x2) (splitJV' x3) (splitJV' x4) (splitJV' x5) + + dynFun <- toSXFun "dynamics" dynamicsFunction + + dynamicsStageFun <- toMXFun "dynamicsStageFunction" (toDynamicsStage callInterpolate cijs dynFun) + callDynamicsStageFun <- fmap call (expandMXFun dynamicsStageFun) + + let nlp :: Nlp (CollTraj x z u p n deg) (JV fp) (CollOcpConstraints x r c h n deg) MX + nlp = Nlp { + nlpFG = + getFg taus + (bcFun :: SXFun ( J (JV x) + :*: J (JV x) + :*: J (JV q) + :*: J (JV p) + :*: J (JV fp) + :*: J (JV Id) + ) + (J (JV c)) + ) + (mayerFun :: SXFun ( J (JV Id) + :*: J (JV x) + :*: J (JV x) + :*: J (JV q) + :*: J (JV p) + :*: J (JV fp) + ) + (J (JV Id)) + ) + (call lagrangeStageFunMX) + (call quadratureStageFunMX) + (call pathCStageFunMX) + (callDynamicsStageFun) + , nlpBX = cat (ocpPhaseBx ocpInputs) + , nlpBG = cat (ocpPhaseBg ocpInputs) + , nlpX0 = guess :: J (CollTraj x z u p n deg) (Vector Double) + , nlpP = catJV (ocpFixedP ocpInputs) + , nlpLamX0 = Nothing + , nlpLamG0 = Nothing + , nlpScaleF = ocpObjScale ocp + , nlpScaleX = Just $ cat $ fillCollTraj + (fromMaybe (fill 1) (ocpXScale ocp)) + (fromMaybe (fill 1) (ocpZScale ocp)) + (fromMaybe (fill 1) (ocpUScale ocp)) + (fromMaybe (fill 1) (ocpPScale ocp)) + (fromMaybe 1 (ocpTScale ocp)) + + , nlpScaleG = Just $ cat $ fillCollConstraints + (fromMaybe (fill 1) (ocpXScale ocp)) + (fromMaybe (fill 1) (ocpResidualScale ocp)) + (fromMaybe (fill 1) (ocpBcScale ocp)) + (fromMaybe (fill 1) (ocpPathCScale ocp)) + } + + -- callbacks and quadrature outputs + lagrangeStageFunFullMX <- toMXFun "lagrangeStageFunFull" lagrangeStageFun + quadratureStageFunFullMX <- toMXFun "quadratureStageFunFull" quadratureStageFun + quadratureOutStageFunFullMX <- toMXFun "quadratureOutStageFunFull" quadratureOutStageFun + + outputFun <- toMXFun "stageOutputs" $ outputFunction callInterpolate cijs taus dynFun + genericQuadraturesFun <- toMXFun "generic quadratures" $ + genericQuadraturesFunction callInterpolateScalar cijs n + + let (getHellaOutputs, getPlotPoints, getOutputs) = toCallbacks n roots taus outputFun pathCStageFunMX lagrangeStageFunFullMX quadratureStageFunFullMX quadratureOutStageFunFullMX quadPlotFunSX + + evalQuadratures :: Vec n (Vec deg Double) -> Double -> IO Double + evalQuadratures qs' tf' = do + let d2d :: Double -> J (JV Id) DMatrix + d2d = realToFrac + qs :: Vec n (J (JVec deg (JV Id)) DMatrix) + qs = fmap (cat . JVec . fmap d2d) qs' + tf :: J (JV Id) DMatrix + tf = realToFrac tf' + evalq :: J (JVec deg (JV Id)) DMatrix -> IO (J (JV Id) DMatrix) + evalq q = eval genericQuadraturesFun (q :*: tf) + stageIntegrals' <- T.mapM evalq qs :: IO (Vec n (J (JV Id) DMatrix)) + let stageIntegrals = fmap (unId . splitJV . d2v) stageIntegrals' :: Vec n Double + return (F.sum stageIntegrals) + + nlpConstraints <- toMXFun "nlp_constraints" (\(x:*:p) -> snd (nlpFG nlp x p)) + let evalConstraints x p = do + g <- eval nlpConstraints (v2d x :*: v2d p) + return (d2v g) + + return $ CollProblem { cpNlp = nlp + , cpOcp = ocp + , cpPlotPoints = getPlotPoints + , cpHellaOutputs = getHellaOutputs + , cpConstraints = evalConstraints + , cpOutputs = getOutputs + , cpTaus = taus + , cpRoots = roots + , cpEvalQuadratures = evalQuadratures + , cpMetaProxy = MetaProxy + } + + +toCallbacks :: + forall x z u p fp r o h q qo po n deg + . ( Vectorize x, Vectorize z, Vectorize u, Vectorize p + , Vectorize o, Vectorize h, Vectorize r, Vectorize q + , Vectorize po, Vectorize qo + , Vectorize fp + , Dim n, Dim deg + ) + => Int + -> QuadratureRoots + -> Vec deg Double + -> MXFun ( J (CollStage (JV x) (JV z) (JV u) deg) + :*: J (JV p) + :*: J (JV fp) + :*: J (JV Id) + :*: J (JV Id) + ) + ( J (JVec deg (JV r)) + :*: J (JVec deg (JV x)) + :*: J (JVec deg (JV o)) + :*: J (JV x) + ) + -> MXFun (PathCStageIn (JV x) (JV z) (JV u) (JV p) (JV fp) deg) (J (JVec deg (JV h))) + -> MXFun (QuadratureStageIn (JV x) (JV z) (JV u) (JV p) (JV fp) deg) (QuadratureStageOut (JV Id) deg) + -> MXFun (QuadratureStageIn (JV x) (JV z) (JV u) (JV p) (JV fp) deg) (QuadratureStageOut (JV q) deg) + -> MXFun (QuadratureStageIn (JV x) (JV z) (JV u) (JV p) (JV fp) deg) (QuadratureStageOut (JV qo) deg) + -> SXFun (QuadraturePlottingIn (JV x) (JV z) (JV u) (JV p) (JV o) (JV q) (JV qo) (JV fp)) (J (JV po)) + -> ( J (CollTraj x z u p n deg) (Vector Double) + -> J (JV fp) (Vector Double) + -> IO ( DynPlotPoints Double + , Vec n (StageOutputs x o h q qo po deg Double) + , Quadratures q qo Double + ) + , J (CollTraj x z u p n deg) (Vector Double) + -> J (JV fp) (Vector Double) + -> IO (DynPlotPoints Double) + , J (CollTraj x z u p n deg) (Vector Double) + -> J (JV fp) (Vector Double) + -> IO (Vec n (StageOutputs x o h q qo po deg Double)) + ) +toCallbacks n roots taus outputFun pathStageConFun lagQuadFun quadFun quadOutFun quadPlotFun = + (getHellaOutputs, getPlotPoints, getOutputs) + where + -- prepare callbacks + f :: J (JV o) DMatrix -> J (JV x) DMatrix -> J (JV h) DMatrix -> J (JV po) DMatrix + -> Quadratures q qo Double -> Quadratures q qo Double + -> ( J (JV o) (Vector Double), J (JV x) (Vector Double), J (JV h) (Vector Double) + , J (JV po) (Vector Double) + , Quadratures q qo Double, Quadratures q qo Double + ) + f o' x' h' po' q q' = (d2v o', d2v x', d2v h', d2v po', q, q') + + callOutputFun :: (J (JV x) DMatrix, J (JV x) DMatrix) + -> J (JV p) (Vector Double) + -> J (JV fp) (Vector Double) + -> J (JV Id) (Vector Double) + -> J (JV Id) DMatrix + -> Quadratures q qo Double + -> ( J (CollStage (JV x) (JV z) (JV u) deg) (Vector Double) + , J (JV Id) (Vector Double) + ) + -> IO ( StageOutputs x o h q qo po deg Double + , Quadratures q qo Double + ) + callOutputFun (x0,xF) p fp h tf previousQuadratures (stage, k) = do + let p' = v2d p + fp' = v2d fp + stage' = v2d stage + (_ :*: xdot :*: out :*: xnext) <- + eval outputFun $ stage' :*: p' :*: fp' :*: (v2d h) :*: (v2d k) + + let stageTimes :: Vec deg (J (JV Id) DMatrix) + stageTimes = fmap (\tau -> t0 + realToFrac tau * h') taus + where + t0 = h' * v2d k + stageTimes' = cat (JVec stageTimes) + h' = v2d h + pathCStageIn = PathCStageIn stage' p' fp' stageTimes' h' + quadratureStageIn = QuadratureStageIn stage' p' fp' stageTimes' h' + hs <- eval pathStageConFun pathCStageIn + QuadratureStageOut lagrQdots lagrQs lagrQNext <- eval lagQuadFun quadratureStageIn + QuadratureStageOut userQdots userQs userQNext <- eval quadFun quadratureStageIn + QuadratureStageOut outQdots outQs outQNext <- eval quadOutFun quadratureStageIn + + let outs0 = unJVec (split out) :: Vec deg (J (JV o) DMatrix) + xdots0 = unJVec (split xdot) :: Vec deg (J (JV x) DMatrix) + hs0 = unJVec (split hs) :: Vec deg (J (JV h) DMatrix) + lagrQs0 = fmap (unId . splitJV . d2v) $ unJVec (split lagrQs) :: Vec deg Double + userQs0 = fmap (splitJV . d2v) $ unJVec (split userQs) :: Vec deg (q Double) + outQs0 = fmap (splitJV . d2v) $ unJVec (split outQs) :: Vec deg (qo Double) + lagrQdots0 = fmap (unId . splitJV . d2v) $ unJVec (split lagrQdots) :: Vec deg Double + userQdots0 = fmap (splitJV . d2v) $ unJVec (split userQdots) :: Vec deg (q Double) + outQdots0 = fmap (splitJV . d2v) $ unJVec (split outQdots) :: Vec deg (qo Double) + qdots = TV.tvzipWith3 Quadratures lagrQdots0 userQdots0 outQdots0 + qs = fmap (previousQuadratures ^+^) $ TV.tvzipWith3 Quadratures lagrQs0 userQs0 outQs0 + + nextQuadratures = + Quadratures + { qLagrange = unId (splitJV (d2v lagrQNext)) + , qUser = splitJV (d2v userQNext) + , qOutputs = splitJV (d2v outQNext) + } ^+^ previousQuadratures + + let quadPlotInputs :: + Vec deg + (QuadraturePlottingIn (JV x) (JV z) (JV u) (JV p) (JV o) (JV q) (JV qo) (JV fp) DMatrix) + quadPlotInputs = + toQuadPlotIn <$> xs <*> zs <*> us <*> outs0 <*> qUsers <*> qOuts <*> stageTimes + qUsers = fmap (v2d . catJV . qUser) qs + qOuts = fmap (v2d . catJV . qOutputs) qs + (xs,zs,us) = TV.tvunzip3 $ fmap (toXzu . split) (unJVec (split xzus)) + where + toXzu (CollPoint x z u) = (x, z, u) + CollStage _ xzus = split stage' + toQuadPlotIn x z u o q qo t = QuadraturePlottingIn x0 xF x z u p' o q qo fp' t tf + + pos <- T.mapM (eval quadPlotFun) quadPlotInputs + + let stageOutputs = + StageOutputs + { soVec = TV.tvzipWith6 f outs0 xdots0 hs0 pos qs qdots + , soXNext = d2v xnext + , soQNext = nextQuadratures + } + + return (stageOutputs, nextQuadratures) + + mapOutputFun :: J (CollTraj x z u p n deg) (Vector Double) + -> J (JV fp) (Vector Double) + -> IO ( Vec n (StageOutputs x o h q qo po deg Double) + , Quadratures q qo Double + ) + mapOutputFun ct fp = do + let CollTraj tf p stages xF = split ct + h = catJV $ Id (tf' / fromIntegral n) + where + Id tf' = splitJV tf + + vstages = unJVec (split stages) + :: Vec n (J (CollStage (JV x) (JV z) (JV u) deg) (Vector Double)) + ks :: Vec n (J (JV Id) (Vector Double)) + ks = TV.mkVec' $ map (catJV . Id . realToFrac) (take n [(0::Int)..]) + + CollStage x0 _ = split (TV.tvhead vstages) + quadratures0 :: Quadratures q qo Double + quadratures0 = fill 0 + mapAccumM (callOutputFun (v2d x0, v2d xF) p fp h (v2d tf)) quadratures0 (TV.tvzip vstages ks) + + getHellaOutputs :: + J (CollTraj x z u p n deg) (Vector Double) + -> J (JV fp) (Vector Double) + -> IO ( DynPlotPoints Double + , Vec n (StageOutputs x o h q qo po deg Double) + , Quadratures q qo Double + ) + getHellaOutputs traj fp = do + (outputs, quadratures) <- mapOutputFun traj fp + return (dynPlotPoints roots (split traj) outputs, outputs, quadratures) + + getPlotPoints :: J (CollTraj x z u p n deg) (Vector Double) + -> J (JV fp) (Vector Double) + -> IO (DynPlotPoints Double) + getPlotPoints traj fp = do + (dpp, _, _) <- getHellaOutputs traj fp + return dpp + + getOutputs :: J (CollTraj x z u p n deg) (Vector Double) + -> J (JV fp) (Vector Double) + -> IO (Vec n (StageOutputs x o h q qo po deg Double)) + getOutputs traj fp = do + (outputs, _) <- mapOutputFun traj fp + return outputs + + +getFg :: + forall x z u p r c h q fp n deg . + ( Dim deg, Dim n + , Vectorize x, Vectorize z, Vectorize u, Vectorize p + , Vectorize r, Vectorize c, Vectorize h, Vectorize q, Vectorize fp + ) + -- taus + => Vec deg Double + -- bcFun + -> SXFun ( J (JV x) + :*: J (JV x) + :*: J (JV q) + :*: J (JV p) + :*: J (JV fp) + :*: J (JV Id) + ) + (J (JV c)) + -- mayerFun + -> SXFun ( J (JV Id) + :*: J (JV x) + :*: J (JV x) + :*: J (JV q) + :*: J (JV p) + :*: J (JV fp) + ) + (J (JV Id)) + -- lagQuadFun + -> (QuadratureStageIn (JV x) (JV z) (JV u) (JV p) (JV fp) deg MX -> J (JV Id) MX) + -- quadFun + -> (QuadratureStageIn (JV x) (JV z) (JV u) (JV p) (JV fp) deg MX -> J (JV q) MX) + -- pathCStageFun + -> (PathCStageIn (JV x) (JV z) (JV u) (JV p) (JV fp) deg MX -> J (JVec deg (JV h)) MX) + -- stageFun + -> ( ( J (JV x) + :*: J (JVec deg (JTuple (JV x) (JV z))) + :*: J (JVec deg (JV u)) + :*: J (JV Id) + :*: J (JV p) + :*: J (JV fp) + :*: J (JVec deg (JV Id)) + ) MX + -> ( J (JVec deg (JV r)) + :*: J (JV x) + ) MX + ) + -- collTraj + -> J (CollTraj x z u p n deg) MX + -- parameter + -> J (JV fp) MX + -- (objective, constraints) + -> (J (JV Id) MX, J (CollOcpConstraints x r c h n deg) MX) +getFg taus bcFun mayerFun lagQuadFun quadFun pathCStageFun dynamicsStageFun collTraj fixedParm = (obj, cat g) + where + -- split up the design vars + CollTraj tf parm stages' xf = split collTraj + stages = unJVec (split stages') :: Vec n (J (CollStage (JV x) (JV z) (JV u) deg) MX) + spstages = fmap split stages :: Vec n (CollStage (JV x) (JV z) (JV u) deg MX) + + obj = objLagrange + objMayer + + objMayer = call mayerFun (tf :*: x0 :*: xf :*: finalQuadratures :*: parm :*: fixedParm) + + objLagrange :: J (JV Id) MX + objLagrange = F.sum $ TV.tvzipWith (oneQuadStage lagQuadFun) stages times' + + finalQuadratures :: J (JV q) MX + finalQuadratures = F.sum $ TV.tvzipWith (oneQuadStage quadFun) stages times' + + oneQuadStage :: + View qOrSomething + => (QuadratureStageIn (JV x) (JV z) (JV u) (JV p) (JV fp) deg MX -> J qOrSomething MX) + -> J (CollStage (JV x) (JV z) (JV u) deg) MX + -> J (JVec deg (JV Id)) MX + -> J qOrSomething MX + oneQuadStage qfun collStage stageTimes = qfun qInputs + where + qInputs :: QuadratureStageIn (JV x) (JV z) (JV u) (JV p) (JV fp) deg MX + qInputs = QuadratureStageIn collStage parm fixedParm stageTimes dt + + -- timestep + dt = tf / fromIntegral n + n = reflectDim (Proxy :: Proxy n) + + -- times at each collocation point + times :: Vec n (Vec deg (J (JV Id) MX)) + times = fmap snd $ timesFromTaus 0 (fmap realToFrac taus) dt + + times' :: Vec n (J (JVec deg (JV Id)) MX) + times' = fmap (cat . JVec) times + + -- initial point at each stage + x0s :: Vec n (J (JV x) MX) + x0s = fmap (\(CollStage x0' _) -> x0') spstages + + -- final point at each stage (for matching constraint) + xfs :: Vec n (J (JV x) MX) + xfs = TV.tvshiftl x0s xf + + x0 = (\(CollStage x0' _) -> x0') (TV.tvhead spstages) + g = CollOcpConstraints + { coCollPoints = cat $ JVec dcs + , coContinuity = cat $ JVec integratorMatchingConstraints + , coPathC = cat $ JVec hs + , coBc = call bcFun (x0 :*: xf :*: finalQuadratures :*: parm :*: fixedParm :*: tf) + } + + integratorMatchingConstraints :: Vec n (J (JV x) MX) -- THIS SHOULD BE A NONLINEAR FUNCTION + integratorMatchingConstraints = vzipWith (-) interpolatedXs xfs + + dcs :: Vec n (J (JVec deg (JV r)) MX) + hs :: Vec n (J (JVec deg (JV h)) MX) + interpolatedXs :: Vec n (J (JV x) MX) + (dcs, hs, interpolatedXs) = TV.tvunzip3 $ fmap fff $ TV.tvzip spstages times' + fff :: (CollStage (JV x) (JV z) (JV u) deg MX, J (JVec deg (JV Id)) MX) -> + (J (JVec deg (JV r)) MX, J (JVec deg (JV h)) MX, J (JV x) MX) + fff (CollStage x0' xzus, stageTimes) = (dc, stageHs, interpolatedX') + where + -- todo: could share xdot here instead of embedding in pathc and dynamics + dc :*: interpolatedX' = + dynamicsStageFun (x0' :*: xzs :*: us :*: dt :*: parm :*: fixedParm :*: stageTimes) + + -- todo: don't split/cat this + pathCStageIn = PathCStageIn (cat (CollStage x0 xzus)) parm fixedParm stageTimes dt + stageHs = pathCStageFun pathCStageIn + + xzs = cat (JVec xzs') :: J (JVec deg (JTuple (JV x) (JV z))) MX + us = cat (JVec us') :: J (JVec deg (JV u)) MX + (xzs', us') = TV.tvunzip $ fmap toTuple $ unJVec (split xzus) + toTuple xzu = (cat (JTuple x z), u) + where + CollPoint x z u = split xzu + + +ocpPhaseBx :: forall x z u p c h fp n deg . + ( Dim n, Dim deg + , Vectorize x, Vectorize z, Vectorize u, Vectorize p + ) + => OcpPhaseInputs x z u p c h fp + -> CollTraj x z u p n deg (Vector Bounds) +ocpPhaseBx ocpInputs = + fillCollTraj' + (fill (Nothing, Nothing)) + (ocpXbnd ocpInputs) + (ocpZbnd ocpInputs) + (ocpUbnd ocpInputs) + (ocpPbnd ocpInputs) + (ocpTbnd ocpInputs) + +ocpPhaseBg :: forall x z u p r c h fp n deg . + ( Dim n, Dim deg + , Vectorize x, Vectorize r, Vectorize c, Vectorize h + ) + => OcpPhaseInputs x z u p c h fp + -> CollOcpConstraints x r c h n deg (Vector Bounds) +ocpPhaseBg ocpInputs = + CollOcpConstraints + { coCollPoints = jreplicate (jfill (Just 0, Just 0)) -- dae residual constraint + , coContinuity = jreplicate (jfill (Just 0, Just 0)) -- continuity constraint + , coPathC = jreplicate (jreplicate hbnds) + , coBc = catJV (ocpBcBnds ocpInputs) + } + where + hbnds :: J (JV h) (Vector Bounds) + hbnds = catJV (ocpPathCBnds ocpInputs) + +toQuadratureFun :: + forall x z u p fp q deg + . ( View q, View x, View z, View u, Dim deg + ) + => Int + -> Vec (TV.Succ deg) (Vec (TV.Succ deg) Double) + -> (J q MX -> Vec deg (J q MX) -> J q MX) + -> (QuadratureIn x z u p fp MX -> J q MX) + -> QuadratureStageIn x z u p fp deg MX + -> (Vec deg (J q MX), Vec deg (J q MX), J q MX) +toQuadratureFun n cijs interpolate' evalQuadDeriv (QuadratureStageIn collStage p fp stageTimes' h) = + (qdots, qs, qnext) + where + CollStage x0 xzus' = split collStage + xzus = fmap split (unJVec (split xzus')) :: Vec deg (CollPoint x z u MX) + tf = h * fromIntegral n + + xs :: Vec deg (J x MX) + xs = fmap (\(CollPoint x _ _) -> x) xzus + + -- state derivatives, maybe these could be useful as outputs + xdots :: Vec deg (J x MX) + xdots = fmap (`M.vs` (1/h)) $ interpolateXDots cijs (x0 TV.<| xs) + + quadratureIns :: Vec deg (QuadratureIn x z u p fp MX) + quadratureIns = TV.tvzipWith3 (\x' (CollPoint x z u) t -> QuadratureIn x' x z u p fp t tf) + xdots xzus stageTimes + + qdots :: Vec deg (J q MX) + qdots = fmap evalQuadDeriv quadratureIns + + stageTimes :: Vec deg (J (JV Id) MX) + stageTimes = unJVec (split stageTimes') + + qnext :: J q MX + qnext = interpolate' (0 :: J q MX) qs + + qs = fmap timesH qsOverH + where + timesH q = M.uncol $ M.ms (M.col q) h + + qsOverH :: Vec deg (J q MX) + qsOverH = cijInvFr !* qdots + + cijs' :: Vec deg (Vec deg Double) + cijs' = TV.tvtail $ fmap TV.tvtail cijs + + cijMat :: Mat.Matrix Double + cijMat = Mat.fromLists $ F.toList $ fmap F.toList cijs' + + cijInv' :: Mat.Matrix Double + cijInv' = Mat.inv cijMat + + cijInv :: Vec deg (Vec deg Double) + cijInv = TV.mkVec' (map TV.mkVec' (Mat.toLists cijInv')) + + cijInvFr :: Vec deg (Vec deg (J q MX)) + cijInvFr = fmap (fmap realToFrac) cijInv + + +toPathCFun :: + forall x z u p fp h deg + . ( View x, View z, View u, View h, Dim deg + ) + => Vec (TV.Succ deg) (Vec (TV.Succ deg) Double) + -> (PathCIn x z u p fp MX -> J h MX) + -> PathCStageIn x z u p fp deg MX + -> Vec deg (J h MX) +toPathCFun cijs evalPathC (PathCStageIn collStage p fp stageTimes' h) = hs + where + CollStage x0 xzus' = split collStage + xzus = fmap split (unJVec (split xzus')) :: Vec deg (CollPoint x z u MX) + + xs :: Vec deg (J x MX) + xs = fmap (\(CollPoint x _ _) -> x) xzus + + -- state derivatives, maybe these could be useful as outputs + xdots :: Vec deg (J x MX) + xdots = fmap (`M.vs` (1/h)) $ interpolateXDots cijs (x0 TV.<| xs) + + pathCIns :: Vec deg (PathCIn x z u p fp MX) + pathCIns = TV.tvzipWith3 (\x' (CollPoint x z u) t -> PathCIn x' x z u p fp t) + xdots xzus stageTimes + + hs :: Vec deg (J h MX) + hs = fmap evalPathC pathCIns + + stageTimes :: Vec deg (J (JV Id) MX) + stageTimes = unJVec (split stageTimes') + + +-- todo: merging this with evaluateQuadraturesFunction would reduce duplication, +-- but could be inefficient +genericQuadraturesFunction :: + forall deg + . Dim deg + => (J (JV Id) MX -> Vec deg (J (JV Id) MX) -> J (JV Id) MX) + -> Vec (TV.Succ deg) (Vec (TV.Succ deg) Double) + -> Int + -> (J (JVec deg (JV Id)) :*: J (JV Id)) MX + -> J (JV Id) MX +genericQuadraturesFunction interpolate' cijs' n (qdots' :*: tf) = + dt * qnext + where + dt = tf / fromIntegral n + + qdots :: Vec deg (J (JV Id) MX) + qdots = unJVec $ split qdots' + + qnext :: J (JV Id) MX + qnext = interpolate' 0 qs + + qs = cijInvFr !* qdots + + cijs :: Vec deg (Vec deg Double) + cijs = TV.tvtail $ fmap TV.tvtail cijs' + + cijMat :: Mat.Matrix Double + cijMat = Mat.fromLists $ F.toList $ fmap F.toList cijs + + cijInv' :: Mat.Matrix Double + cijInv' = Mat.inv cijMat + + cijInv :: Vec deg (Vec deg Double) + cijInv = TV.mkVec' (map TV.mkVec' (Mat.toLists cijInv')) + + cijInvFr :: Vec deg (Vec deg (J (JV Id) MX)) + cijInvFr = fmap (fmap realToFrac) cijInv + + +-- todo: code duplication +dot :: forall x deg a b. (Fractional (J x a), Real b, Dim deg) => Vec deg b -> Vec deg (J x a) -> J x a +dot cks xs = F.sum $ TV.unVec elemwise + where + elemwise :: Vec deg (J x a) + elemwise = TV.tvzipWith smul cks xs + + smul :: b -> J x a -> J x a + smul x y = realToFrac x * y + + +-- todo: code duplication +interpolateXDots' :: (Real b, Fractional (J x a), Dim deg) => Vec deg (Vec deg b) -> Vec deg (J x a) -> Vec deg (J x a) +interpolateXDots' cjks xs = fmap (`dot` xs) cjks + +interpolateXDots :: + (Real b, Dim deg, Fractional (J x a)) => + Vec (TV.Succ deg) (Vec (TV.Succ deg) b) + -> Vec (TV.Succ deg) (J x a) + -> Vec deg (J x a) +interpolateXDots cjks xs = TV.tvtail $ interpolateXDots' cjks xs + + +-- return dynamics constraints and interpolated state +toDynamicsStage :: + forall x z u p fp r o deg . (Dim deg, View x, View z, View u, View p, View fp, View r, View o) + => (J x MX -> Vec deg (J x MX) -> J x MX) + -> Vec (TV.Succ deg) (Vec (TV.Succ deg) Double) + -> SXFun (DaeIn x z u p fp) (DaeOut r o) + -> (J x :*: J (JVec deg (JTuple x z)) :*: J (JVec deg u) :*: J (JV Id) :*: J p :*: J fp :*: J (JVec deg (JV Id))) MX + -> (J (JVec deg r) :*: J x) MX +toDynamicsStage interpolate' cijs dynFun (x0 :*: xzs' :*: us' :*: h :*: p :*: fp :*: stageTimes') = + cat (JVec dynConstrs) :*: xnext + where + xzs = fmap split (unJVec (split xzs')) :: Vec deg (JTuple x z MX) + us = unJVec (split us') :: Vec deg (J u MX) + + -- interpolated final state + xnext :: J x MX + xnext = interpolate' x0 xs + + stageTimes = unJVec $ split stageTimes' + + -- dae constraints (dynamics) + dynConstrs :: Vec deg (J r MX) + (dynConstrs, _) = TV.tvunzip $ TV.tvzipWith4 applyDae xdots xzs us stageTimes + + applyDae :: J x MX -> JTuple x z MX -> J u MX -> J (JV Id) MX -> (J r MX, J o MX) + applyDae x' (JTuple x z) u t = (r, o) + where + DaeOut r o = call dynFun (DaeIn t p fp x' collPoint) + collPoint = cat (CollPoint x z u) + + -- state derivatives, maybe these could be useful as outputs + xdots :: Vec deg (J x MX) + xdots = fmap (`M.vs` (1/h)) $ interpolateXDots cijs (x0 TV.<| xs) + + xs :: Vec deg (J x MX) + xs = fmap (\(JTuple x _) -> x) xzs + + +-- outputs +outputFunction :: + forall x z u p fp r o deg . (Dim deg, View x, View z, View u, View p, View fp, View r, View o) + => (J x MX -> Vec deg (J x MX) -> J x MX) + -> Vec (TV.Succ deg) (Vec (TV.Succ deg) Double) -> Vec deg Double + -> SXFun (DaeIn x z u p fp) (DaeOut r o) + -> (J (CollStage x z u deg) :*: J p :*: J fp :*: J (JV Id) :*: J (JV Id)) MX + -> (J (JVec deg r) :*: J (JVec deg x) :*: J (JVec deg o) :*: J x) MX +outputFunction callInterpolate cijs taus dynFun (collStage :*: p :*: fp :*: h :*: k) = + cat (JVec dynConstrs) :*: cat (JVec xdots) :*: cat (JVec outputs) :*: xnext + where + xzus = unJVec (split xzus') :: Vec deg (J (CollPoint x z u) MX) + CollStage x0 xzus' = split collStage + -- times at each collocation point + stageTimes :: Vec deg (J (JV Id) MX) + stageTimes = fmap (\tau -> t0 + realToFrac tau * h) taus + t0 = k*h + + xnext = callInterpolate x0 xs + + -- dae constraints (dynamics) + dynConstrs :: Vec deg (J r MX) + outputs :: Vec deg (J o MX) + (dynConstrs, outputs) = TV.tvunzip $ TV.tvzipWith3 applyDae xdots xzus stageTimes + + applyDae :: J x MX -> J (CollPoint x z u) MX -> J (JV Id) MX -> (J r MX, J o MX) + applyDae x' xzu t = (r, o) + where + DaeOut r o = call dynFun (DaeIn t p fp x' xzu) + + -- state derivatives, maybe these could be useful as outputs + xdots :: Vec deg (J x MX) + xdots = fmap (`M.vs` (1/h)) $ interpolateXDots cijs (x0 TV.<| xs) + + xs :: Vec deg (J x MX) + xs = fmap ((\(CollPoint x _ _) -> x) . split) xzus + + + +-- | make an initial guess +makeGuess :: + forall x z u p deg n . + ( Dim n, Dim deg + , Vectorize x, Vectorize z, Vectorize u, Vectorize p + ) + => QuadratureRoots + -> Double -> (Double -> x Double) -> (Double -> z Double) -> (Double -> u Double) + -> p Double + -> CollTraj x z u p n deg (Vector Double) +makeGuess quadratureRoots tf guessX guessZ guessU parm = + CollTraj (jfill tf) (catJV parm) guesses (catJV (guessX tf)) + where + -- timestep + dt = tf / fromIntegral n + n = vlength (Proxy :: Proxy (Vec n)) + + -- initial time at each collocation stage + t0s :: Vec n Double + t0s = TV.mkVec' $ take n [dt * fromIntegral k | k <- [(0::Int)..]] + + -- times at each collocation point + times :: Vec n (Double, Vec deg Double) + times = fmap (\t0 -> (t0, fmap (\tau -> t0 + tau*dt) taus)) t0s + + mkGuess' :: (Double, Vec deg Double) -> CollStage (JV x) (JV z) (JV u) deg (Vector Double) + mkGuess' (t,ts) = + CollStage (catJV (guessX t)) $ + cat $ JVec $ fmap (\t' -> cat (CollPoint (catJV (guessX t')) (catJV (guessZ t')) (catJV (guessU t')))) ts + + guesses :: J (JVec n (CollStage (JV x) (JV z) (JV u) deg)) (Vector Double) + guesses = cat $ JVec $ fmap (cat . mkGuess') times + + -- the collocation points + taus :: Vec deg Double + taus = mkTaus quadratureRoots + + +-- | make an initial guess +makeGuessSim :: + forall x z u p deg n . + ( Dim n, Dim deg + , Vectorize x, Vectorize z, Vectorize u, Vectorize p + , Additive x + ) + => QuadratureRoots + -> Double + -> x Double + -> (Double -> x Double -> u Double -> x Double) + -> (Double -> x Double -> u Double) + -> p Double + -> CollTraj x z u p n deg (Vector Double) +makeGuessSim quadratureRoots tf x00 ode guessU p = + CollTraj (jfill tf) (catJV p) (cat (JVec stages)) (catJV xf) + where + -- timestep + dt = tf / fromIntegral n + n = vlength (Proxy :: Proxy (Vec n)) + + -- initial time at each collocation stage + t0s :: Vec n Double + t0s = TV.mkVec' $ take n [dt * fromIntegral k | k <- [(0::Int)..]] + + xf :: x Double + stages :: Vec n (J (CollStage (JV x) (JV z) (JV u) deg) (Vector Double)) + (xf, stages) = T.mapAccumL stageGuess x00 t0s + + stageGuess :: x Double -> Double + -> (x Double, J (CollStage (JV x) (JV z) (JV u) deg) (Vector Double)) + stageGuess x0 t0 = (fst (integrate 1), cat (CollStage (catJV x0) points)) + where + points = cat $ JVec $ fmap (toCollPoint . integrate) taus + f :: Double -> x Double -> x Double + f t x = ode t x u + where + u = guessU t x + toCollPoint (x,u) = cat $ CollPoint (catJV x) (catJV (fill 0 :: z Double)) (catJV u) + integrate localTau = (x, u) + where + t = localTau * dt + x = rk45 f (InitialTime t0) (TimeStep t) x0 + u = guessU t x + + -- the collocation points + taus :: Vec deg Double + taus = mkTaus quadratureRoots + + +-- http://stackoverflow.com/questions/11652809/how-to-implement-mapaccumm +-- thanks rconner +mapAccumM :: (Monad m, Functor m, T.Traversable t) => (a -> b -> m (c, a)) -> a -> t b -> m (t c, a) +mapAccumM f = flip (runStateT . (T.traverse (StateT . (flip f)))) diff --git a/src/Dyno/DirectCollocation/FormulateCov.hs b/src/Dyno/DirectCollocation/FormulateCov.hs new file mode 100644 --- /dev/null +++ b/src/Dyno/DirectCollocation/FormulateCov.hs @@ -0,0 +1,300 @@ +{-# OPTIONS_GHC -Wall #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE PolyKinds #-} + +module Dyno.DirectCollocation.FormulateCov + ( CollCovProblem(..) + , CovTraj(..) + , makeCollCovProblem + ) where + +import Data.Maybe ( fromMaybe ) +import Data.Proxy ( Proxy(..) ) +import Data.Vector ( Vector ) +import qualified Data.Foldable as F +import Linear.V + +import Casadi.DMatrix ( DMatrix ) +import Casadi.MX ( MX ) + +import Dyno.View.View ( View(..), J, jfill, v2d, d2v ) +import Dyno.View.Cov ( Cov ) +import Dyno.View.JV ( JV, catJV, catJV', splitJV' ) +import Dyno.View.HList ( (:*:)(..) ) +import Dyno.View.Fun +import Dyno.View.JVec( JVec(..), jreplicate ) +import Dyno.Vectorize ( Vectorize(..), Id(..), None(..), fill ) +import Dyno.TypeVecs ( Vec ) +import qualified Dyno.TypeVecs as TV +import Dyno.Nlp ( Nlp(..), Bounds ) +import Dyno.Ocp + +import Dyno.DirectCollocation.Types +import Dyno.DirectCollocation.Dynamic ( DynPlotPoints ) +import Dyno.DirectCollocation.Quadratures ( QuadratureRoots(..), timesFromTaus ) +import Dyno.DirectCollocation.Robust +import Dyno.DirectCollocation.Formulate + +data CollCovProblem ocp n deg sx sw sh shr sc = + CollCovProblem + { ccpNlp :: Nlp + (CollTrajCov sx ocp n deg) + (JV None) + (CollOcpCovConstraints ocp n deg sh shr sc) MX + , ccpPlotPoints :: J (CollTrajCov sx ocp n deg) (Vector Double) -> IO (DynPlotPoints Double) + , ccpOutputs :: + J (CollTrajCov sx ocp n deg) (Vector Double) + -> IO ( Vec n (StageOutputs (X ocp) (O ocp) (H ocp) (Q ocp) (QO ocp) (PO ocp) deg Double) + , Vec n (J (Cov (JV sx)) (Vector Double)) + , J (Cov (JV sx)) (Vector Double) + ) + , ccpSensitivities :: MXFun + (J (CollTraj' ocp n deg)) + (CovarianceSensitivities (JV sx) (JV sw) n) + , ccpCovariances :: MXFun + (J (Cov (JV sx)) :*: J (CollTraj (X ocp) (Z ocp) (U ocp) (P ocp) n deg)) + (J (CovTraj sx n)) + , ccpRoots :: QuadratureRoots + } + + + + +makeCollCovProblem :: + forall ocp x z u p fp r o c h q qo po sx sz sw sr sh shr sc deg n . + ( Dim deg, Dim n, Vectorize x, Vectorize p, Vectorize u, Vectorize z + , Vectorize sr, Vectorize sw, Vectorize sz, Vectorize sx + , Vectorize r, Vectorize o, Vectorize h, Vectorize c, Vectorize q, Vectorize po + , Vectorize qo + , View sh, Vectorize shr, View sc + , x ~ X ocp + , q ~ Q ocp + , qo ~ QO ocp + , h ~ H ocp + , c ~ C ocp + , o ~ O ocp + , r ~ R ocp + , p ~ P ocp + , u ~ U ocp + , z ~ Z ocp + , po ~ PO ocp + , fp ~ None + , None ~ FP ocp + ) + => QuadratureRoots + -> OcpPhase' ocp + -> OcpPhaseInputs x z u p c h fp + -> OcpPhaseWithCov ocp sx sz sw sr sh shr sc + -> J (CollTraj x z u p n deg) (Vector Double) + -> IO (CollCovProblem ocp n deg sx sw sh shr sc) +makeCollCovProblem roots ocp ocpInputs ocpCov guess = do + let -- the collocation points + taus :: Vec deg Double + taus = mkTaus roots + + computeSensitivities <- mkComputeSensitivities roots (ocpCovDae ocpCov) + computeCovariances <- mkComputeCovariances continuousToDiscreetNoiseApprox + (computeSensitivities) (ocpCovSq ocpCov) + + sbcFun <- toSXFun "sbc" $ \(x0:*:x1) -> ocpCovSbc ocpCov x0 x1 + shFun <- toSXFun "sh" $ \(x0:*:x1) -> ocpCovSh ocpCov (splitJV' x0) x1 + mayerFun <- toSXFun "cov mayer" $ \(x0:*:x1:*:x2:*:x3:*:x4) -> + catJV' $ Id $ ocpCovMayer ocpCov (unId (splitJV' x0)) (splitJV' x1) (splitJV' x2) x3 x4 + lagrangeFun <- toSXFun "cov lagrange" $ \(x0:*:x1:*:x2:*:x3) -> + catJV' $ Id $ ocpCovLagrange ocpCov (unId (splitJV' x0)) (splitJV' x1) x2 (unId (splitJV' x3)) + + cp0 <- makeCollProblem roots ocp ocpInputs guess + + robustify <- mkRobustifyFunction (ocpCovProjection ocpCov) (ocpCovRobustifyPathC ocpCov) + + let nlp0 = cpNlp cp0 + gammas' = ocpCovGammas ocpCov :: shr Double + + gammas :: J (JV shr) MX + gammas = catJV' (fmap realToFrac gammas') + + rpathCUb :: shr Bounds + rpathCUb = fill (Nothing, Just 0) + + robustPathCUb :: J (JV shr) (Vector Bounds) + robustPathCUb = catJV rpathCUb + + -- the NLP + fg :: J (CollTrajCov sx ocp n deg) MX + -> J (JV fp) MX + -> (J (JV Id) MX, J (CollOcpCovConstraints ocp n deg sh shr sc) MX) + fg = getFgCov taus + computeCovariances + gammas + (robustify :: (J (JV shr) MX -> J (JV p) MX -> J (JV x) MX -> J (Cov (JV sx)) MX -> J (JV shr) MX)) + (sbcFun :: SXFun (J (Cov (JV sx)) :*: J (Cov (JV sx))) (J sc)) + (shFun :: SXFun (J (JV x) :*: J (Cov (JV sx))) (J sh)) + (lagrangeFun :: SXFun (J (JV Id) :*: J (JV x) :*: J (Cov (JV sx)) :*: J (JV Id)) (J (JV Id))) + (mayerFun :: SXFun (J (JV Id) :*: (J (JV x) :*: (J (JV x) :*: (J (Cov (JV sx)) :*: J (Cov (JV sx)))))) (J (JV Id))) + (nlpFG nlp0) + + computeCovariancesFun' <- toMXFun "compute covariances" (\(x :*: y) -> computeCovariances x y) + -- callbacks + let getPlotPoints :: J (CollTrajCov sx ocp n deg) (Vector Double) -> IO (DynPlotPoints Double) + getPlotPoints collTrajCov = do + let CollTrajCov _ collTraj = split collTrajCov + cpPlotPoints cp0 collTraj (catJV None) + + getOutputs :: J (CollTrajCov sx ocp n deg) (Vector Double) + -> IO ( Vec n (StageOutputs x o h q qo po deg Double) + , Vec n (J (Cov (JV sx)) (Vector Double)) + , J (Cov (JV sx)) (Vector Double) + ) + getOutputs collTrajCov = do + let CollTrajCov p0 collTraj = split collTrajCov + outputs <- (cpOutputs cp0) collTraj (catJV None) + covTraj <- fmap split $ eval computeCovariancesFun' (v2d p0 :*: v2d collTraj) + let covs' = ctAllButLast covTraj + pF = ctLast covTraj + let covs = unJVec (split covs') :: Vec n (J (Cov (JV sx)) DMatrix) + return (outputs, fmap d2v covs, d2v pF) + + nlp = + Nlp + { nlpFG = fg + , nlpBX = cat $ CollTrajCov (ocpCovS0bnd ocpCov) (nlpBX nlp0) + , nlpBG = cat $ CollOcpCovConstraints + { cocNormal = nlpBG nlp0 + , cocCovPathC = jreplicate (ocpCovShBnds ocpCov) + , cocCovRobustPathC = jreplicate robustPathCUb + , cocSbc = ocpCovSbcBnds ocpCov + } + , nlpX0 = cat $ CollTrajCov (jfill 0) (nlpX0 nlp0) + , nlpP = catJV None + , nlpLamX0 = Nothing + , nlpLamG0 = Nothing + , nlpScaleF = ocpObjScale ocp + , nlpScaleX = Just $ cat $ + CollTrajCov (fromMaybe (jfill 1) (ocpCovSScale ocpCov)) $ + cat $ fillCollTraj + (fromMaybe (fill 1) (ocpXScale ocp)) + (fromMaybe (fill 1) (ocpZScale ocp)) + (fromMaybe (fill 1) (ocpUScale ocp)) + (fromMaybe (fill 1) (ocpPScale ocp)) + (fromMaybe 1 (ocpTScale ocp)) + + , nlpScaleG = Just $ cat $ CollOcpCovConstraints + { cocNormal = cat $ fillCollConstraints + (fromMaybe (fill 1) (ocpXScale ocp)) + (fromMaybe (fill 1) (ocpResidualScale ocp)) + (fromMaybe (fill 1) (ocpBcScale ocp)) + (fromMaybe (fill 1) (ocpPathCScale ocp)) + , cocCovPathC = jreplicate (fromMaybe (jfill 1) (ocpCovPathCScale ocpCov)) + , cocCovRobustPathC = jreplicate $ + fromMaybe (jfill 1) $ + fmap catJV (ocpCovRobustPathCScale ocpCov) + , cocSbc = fromMaybe (jfill 1) (ocpCovSbcScale ocpCov) + } + } + computeSensitivitiesFun' <- toMXFun "compute sensitivities" computeSensitivities + return $ CollCovProblem { ccpNlp = nlp + , ccpPlotPoints = getPlotPoints + , ccpOutputs = getOutputs + , ccpSensitivities = computeSensitivitiesFun' + , ccpCovariances = computeCovariancesFun' + , ccpRoots = roots + } + + +getFgCov :: + forall ocp x z u p r c h fp sx sh shr sc n deg . + ( Dim deg, Dim n, Vectorize x, Vectorize z, Vectorize u, Vectorize p + , Vectorize h, Vectorize c, Vectorize r, Vectorize fp + , Vectorize sx, View sc, View sh, Vectorize shr + , X ocp ~ x + , Z ocp ~ z + , U ocp ~ u + , P ocp ~ p + , R ocp ~ r + , C ocp ~ c + , H ocp ~ h + , FP ocp ~ fp + ) + -- taus + => Vec deg Double + -> (J (Cov (JV sx)) MX -> J (CollTraj x z u p n deg) MX -> J (CovTraj sx n) MX) + -- gammas + -> J (JV shr) MX + -- robustify + -> (J (JV shr) MX -> J (JV p) MX -> J (JV x) MX -> J (Cov (JV sx)) MX -> J (JV shr) MX) + -- sbcFun + -> SXFun (J (Cov (JV sx)) :*: J (Cov (JV sx))) (J sc) + -- shFun + -> SXFun (J (JV x) :*: J (Cov (JV sx))) (J sh) + -- lagrangeFun + -> SXFun + (J (JV Id) :*: J (JV x) :*: J (Cov (JV sx)) :*: J (JV Id)) (J (JV Id)) + -- mayerFun + -> SXFun + (J (JV Id) :*: J (JV x) :*: J (JV x) :*: J (Cov (JV sx)) :*: J (Cov (JV sx))) (J (JV Id)) + -> (J (CollTraj' ocp n deg) MX -> J (JV fp) MX -> (J (JV Id) MX, J (CollOcpConstraints' ocp n deg) MX) + ) + -> J (CollTrajCov sx ocp n deg) MX + -> J (JV fp) MX + -> (J (JV Id) MX, J (CollOcpCovConstraints ocp n deg sh shr sc) MX) +getFgCov + taus computeCovariances + gammas robustify sbcFun shFun lagrangeFun mayerFun + normalFG collTrajCov nlpParams = + (obj0 + objectiveLagrangeCov + objectiveMayerCov, cat g) + where + CollTrajCov p0 collTraj = split collTrajCov + (obj0, g0) = normalFG collTraj nlpParams + + g = CollOcpCovConstraints + { cocNormal = g0 + , cocCovPathC = cat (JVec covPathConstraints) + , cocCovRobustPathC = cat (JVec robustifiedPathC) + , cocSbc = call sbcFun (p0 :*: pF) + } + -- split up the design vars + CollTraj tf parm stages' xf = split collTraj + stages = unJVec (split stages') :: Vec n (J (CollStage (JV x) (JV z) (JV u) deg) MX) + spstages = fmap split stages :: Vec n (CollStage (JV x) (JV z) (JV u) deg MX) + + objectiveMayerCov = call mayerFun (tf :*: x0 :*: xf :*: p0 :*: pF) + + -- timestep + dt = tf / fromIntegral n + n = reflectDim (Proxy :: Proxy n) + + -- times at each collocation point + t0s :: Vec n (J (JV Id) MX) + (t0s, _) = TV.tvunzip $ timesFromTaus 0 (fmap realToFrac taus) dt + + -- initial point at each stage + x0s :: Vec n (J (JV x) MX) + x0s = fmap (\(CollStage x0' _) -> x0') spstages + + x0 = (\(CollStage x0' _) -> x0') (TV.tvhead spstages) + +-- sensitivities = call computeSensitivities collTraj + + covs :: Vec n (J (Cov (JV sx)) MX) + covs = unJVec (split covs') + + covs' :: J (JVec n (Cov (JV sx))) MX -- all but last covariance + pF :: J (Cov (JV sx)) MX -- last covariances + CovTraj covs' pF = split (computeCovariances p0 collTraj) + + -- lagrange term + objectiveLagrangeCov = (lagrangeF + lagrange0s) / fromIntegral n + where + lagrangeF = call lagrangeFun (tf :*: xf :*: pF :*: tf) + lagrange0s = + sum $ F.toList $ + TV.tvzipWith3 (\tk xk pk -> call lagrangeFun (tk :*: xk :*: pk :*: tf)) t0s x0s covs + + covPathConstraints :: Vec n (J sh MX) + covPathConstraints = TV.tvzipWith (\xk pk -> call shFun (xk:*:pk)) x0s covs + + robustifiedPathC :: Vec n (J (JV shr) MX) + robustifiedPathC = TV.tvzipWith (robustify gammas parm) x0s covs + diff --git a/src/Dyno/DirectCollocation/Integrate.hs b/src/Dyno/DirectCollocation/Integrate.hs --- a/src/Dyno/DirectCollocation/Integrate.hs +++ b/src/Dyno/DirectCollocation/Integrate.hs @@ -1,8 +1,8 @@ {-# OPTIONS_GHC -Wall #-} -{-# Language ScopedTypeVariables #-} -{-# Language TypeOperators #-} -{-# Language DeriveGeneric #-} -{-# Language FlexibleContexts #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE FlexibleContexts #-} module Dyno.DirectCollocation.Integrate ( withIntegrator @@ -17,10 +17,10 @@ import qualified Data.Foldable as F import Linear.V +import Casadi.SX ( SX ) import Casadi.MX ( MX ) -import Dyno.SXElement ( SXElement, sxSplitJV, sxCatJV ) -import Dyno.View.JV ( JV, splitJV, catJV ) +import Dyno.View.JV ( JV, splitJV, catJV, splitJV', catJV' ) import Dyno.View.Viewable ( Viewable ) import Dyno.View.View ( View(..), J, JNone, JTuple(..), jfill ) import Dyno.View.Fun ( SXFun, call, toSXFun, toMXFun, expandMXFun ) @@ -37,7 +37,7 @@ import Dyno.DirectCollocation.Types ( CollStage(..), CollPoint(..) ) import Dyno.DirectCollocation.Quadratures ( QuadratureRoots, mkTaus, interpolate, timesFromTaus ) - +type Sxe = J (JV Id) SX data IntegratorX x z n deg a = IntegratorX @@ -135,9 +135,6 @@ where CollPoint x z u = split collPoint -type Sxe = SXElement - - withIntegrator :: forall x z u p r deg n b . (Dim n, Dim deg, Vectorize x, Vectorize p, Vectorize u, Vectorize z, Vectorize r) @@ -161,9 +158,9 @@ dynFun <- toSXFun "dynamics" $ dynamicsFunction' $ \x0 x1 x2 x3 x4 x5 -> - let r = dae (sxSplitJV x0) (sxSplitJV x1) (sxSplitJV x2) (sxSplitJV x3) - (sxSplitJV x4) (unId (sxSplitJV x5)) - in sxCatJV r + let r = dae (splitJV' x0) (splitJV' x1) (splitJV' x2) (splitJV' x3) + (splitJV' x4) (unId (splitJV' x5)) + in catJV' r dynStageConFun <- toMXFun "dynamicsStageCon" (dynStageConstraints' cijs taus dynFun) -- let callDynStageConFun = call dynStageConFun diff --git a/src/Dyno/DirectCollocation/Interpolate.hs b/src/Dyno/DirectCollocation/Interpolate.hs --- a/src/Dyno/DirectCollocation/Interpolate.hs +++ b/src/Dyno/DirectCollocation/Interpolate.hs @@ -1,9 +1,9 @@ {-# OPTIONS_GHC -Wall #-} -{-# Language ScopedTypeVariables #-} -{-# Language DeriveFunctor #-} -{-# Language DeriveFoldable #-} -{-# Language DeriveTraversable #-} -{-# Language PolyKinds #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE DeriveFunctor #-} +{-# LANGUAGE DeriveFoldable #-} +{-# LANGUAGE DeriveTraversable #-} +{-# LANGUAGE PolyKinds #-} module Dyno.DirectCollocation.Interpolate ( interpolateTraj @@ -67,7 +67,7 @@ type Point x z u = CollPoint (JV x) (JV z) (JV u) -newtype Times deg a = Times (a, Vec deg a) deriving (Functor, F.Foldable, T.Traversable) +newtype Times deg a = Times (a, Vec deg a) deriving Functor -- | re-discretize a collocation trajectory using the lagrange interpolation polynomials -- from the quadrature scheme diff --git a/src/Dyno/DirectCollocation/Profile.hs b/src/Dyno/DirectCollocation/Profile.hs deleted file mode 100644 --- a/src/Dyno/DirectCollocation/Profile.hs +++ /dev/null @@ -1,69 +0,0 @@ -{-# OPTIONS_GHC -Wall #-} -{-# Language ScopedTypeVariables #-} -{-# Language RankNTypes #-} - -module Dyno.DirectCollocation.Profile - ( ProfileReport(..) - , profile - ) where - -import Data.Proxy ( Proxy(..) ) -import Data.Vector ( Vector ) -import Linear.V ( Dim(..) ) - -import Dyno.View.View ( J ) -import Dyno.Vectorize ( Vectorize ) -import Dyno.Ocp -import Dyno.Solvers ( Solver ) -import Dyno.DirectCollocation.Types ( CollTraj, CollOcpConstraints ) -import Dyno.DirectCollocation.Formulate ( CollProblem(..), makeCollProblem ) -import Dyno.DirectCollocation.Quadratures ( QuadratureRoots ) -import qualified Dyno.TypeVecs as TV -import Dyno.NlpUtils ( solveNlp ) -import Dyno.Nlp ( NlpOut(..) ) - -data ProfileReport = - ProfileReport - { - } - -toProfileReport :: - Either String String - -> NlpOut (CollTraj x z u p n deg) (CollOcpConstraints x r c h n deg) (Vector Double) - -> IO ProfileReport -toProfileReport _ _ = return ProfileReport - -profile :: forall x z u p r o c h q . - ( Vectorize x, Vectorize z, Vectorize u, Vectorize p - , Vectorize r, Vectorize o, Vectorize c, Vectorize h, Vectorize q - ) - => QuadratureRoots - -> OcpPhase x z u p r o c h q - -> (forall deg n . (Dim deg, Dim n) => J (CollTraj x z u p n deg) (Vector Double)) - -> Solver - -> [(Int,Int)] - -> IO [ProfileReport] -profile roots ocp guess solver range = do - let go :: (Int,Int) -> IO ProfileReport - go (n,deg) = - TV.reifyDim n $ \(Proxy :: Proxy n ) -> - TV.reifyDim deg $ \(Proxy :: Proxy deg) -> - profileOne roots ocp (guess :: J (CollTraj x z u p n deg) (Vector Double)) solver - mapM go range - -profileOne :: - forall x z u p r o c h q n deg . - ( Vectorize x, Vectorize z, Vectorize u, Vectorize p - , Vectorize r, Vectorize o, Vectorize c, Vectorize h, Vectorize q - , Dim n, Dim deg - ) - => QuadratureRoots - -> OcpPhase x z u p r o c h q - -> J (CollTraj x z u p n deg) (Vector Double) - -> Solver - -> IO ProfileReport -profileOne roots ocp guess solver = do - cp <- makeCollProblem roots ocp guess - let nlp = cpNlp cp - x <- solveNlp solver nlp Nothing - uncurry toProfileReport x diff --git a/src/Dyno/DirectCollocation/Quadratures.hs b/src/Dyno/DirectCollocation/Quadratures.hs --- a/src/Dyno/DirectCollocation/Quadratures.hs +++ b/src/Dyno/DirectCollocation/Quadratures.hs @@ -1,8 +1,8 @@ {-# OPTIONS_GHC -Wall #-} -{-# Language ScopedTypeVariables #-} -{-# Language DeriveGeneric #-} -{-# Language FlexibleContexts #-} -{-# Language PolyKinds #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE PolyKinds #-} module Dyno.DirectCollocation.Quadratures ( QuadratureRoots(..) @@ -14,6 +14,7 @@ import GHC.Generics ( Generic ) +import Data.Aeson ( FromJSON, ToJSON ) import Data.Proxy ( Proxy(..) ) import qualified Data.Vector as V import qualified Data.Foldable as F @@ -25,19 +26,22 @@ import Dyno.View.View ( View, J ) import Dyno.TypeVecs ( Vec ) +import Dyno.Vectorize ( devectorize ) import qualified Dyno.TypeVecs as TV import Dyno.LagrangePolynomials ( lagrangeXis ) data QuadratureRoots = Legendre | Radau deriving (Show, Eq, Ord, Enum, Generic) instance Binary QuadratureRoots instance Serialize QuadratureRoots +instance ToJSON QuadratureRoots +instance FromJSON QuadratureRoots mkTaus :: forall deg a . (Dim deg, Fractional a) => QuadratureRoots -> Vec deg a mkTaus quadratureRoots = case taus of - Just taus' -> TV.mkVec $ V.map (fromRational . toRational) taus' + Just taus' -> devectorize $ V.map (fromRational . toRational) taus' Nothing -> error "makeTaus: too high degree" where deg = reflectDim (Proxy :: Proxy deg) diff --git a/src/Dyno/DirectCollocation/Robust.hs b/src/Dyno/DirectCollocation/Robust.hs --- a/src/Dyno/DirectCollocation/Robust.hs +++ b/src/Dyno/DirectCollocation/Robust.hs @@ -1,9 +1,9 @@ {-# OPTIONS_GHC -Wall #-} -{-# Language TypeFamilies #-} -{-# Language ScopedTypeVariables #-} -{-# Language TypeOperators #-} -{-# Language DeriveGeneric #-} -{-# Language FlexibleContexts #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE PolyKinds #-} module Dyno.DirectCollocation.Robust @@ -27,8 +27,6 @@ import qualified Dyno.View.Unsafe.M as M ( mkM, blockSplit ) -import Dyno.SXElement ( SXElement, sxSplitJV, sxCatJV ) -import Dyno.Ocp import Dyno.View.View ( View(..), J, JNone(..), JTuple(..), fromDMatrix ) import Dyno.View.JV ( JV, catJV', splitJV' ) import Dyno.View.HList ( (:*:)(..) ) @@ -63,7 +61,7 @@ } deriving (Eq, Show, Generic, Generic1) instance (View xe, View we, Dim n) => Scheme (CovarianceSensitivities xe we n) -type Sxe = SXElement +type Sxe = J (JV Id) SX mkComputeSensitivities :: forall x z u p sx sz sw sr deg n . @@ -87,16 +85,16 @@ errorDynFun <- toSXFun "error dynamics" $ errorDynamicsFunction $ \x0 x1 x2 x3 x4 x5 x6 x7 x8 x9 -> let r = covDae - (sxSplitJV x0) (sxSplitJV x1) (sxSplitJV x2) (sxSplitJV x3) (sxSplitJV x4) - (unId (sxSplitJV x5)) (sxSplitJV x6) (sxSplitJV x7) (sxSplitJV x8) (sxSplitJV x9) - in sxCatJV r + (splitJV' x0) (splitJV' x1) (splitJV' x2) (splitJV' x3) (splitJV' x4) + (unId (splitJV' x5)) (splitJV' x6) (splitJV' x7) (splitJV' x8) (splitJV' x9) + in catJV' r edscf <- toMXFun "errorDynamicsStageCon" (errorDynStageConstraints cijs taus errorDynFun) errorDynStageConFunJac <- toFunJac edscf sensitivityStageFun' <- toMXFun "sensitivity stage function" $ sensitivityStageFunction (call errorDynStageConFunJac) - sensitivityStageFun <- expandMXFun sensitivityStageFun' + let sensitivityStageFun = sensitivityStageFun' let sens :: J (JV Id) MX -> J (JV p) MX -> J (JVec deg (JV Id)) MX @@ -142,28 +140,22 @@ -- todo: calculate by first multiplying all the Fs mkComputeCovariances :: - forall ocp x z u p sx sw n deg . + forall x z u p sx sw n deg . ( Dim deg, Dim n , Vectorize x, Vectorize z, Vectorize u, Vectorize p , Vectorize sx, Vectorize sw - , X ocp ~ x - , Z ocp ~ z - , U ocp ~ u - , P ocp ~ p ) => (M (JV sx) (JV sx) MX -> M (JV sx) (JV sw) MX -> J (Cov (JV sw)) MX -> J (JV Id) MX -> M (JV sx) (JV sx) MX) -> (J (CollTraj x z u p n deg) MX -> CovarianceSensitivities (JV sx) (JV sw) n MX) -> J (Cov (JV sw)) DMatrix - -> IO (J (CollTrajCov sx ocp n deg) MX -> J (CovTraj sx n) MX) + -> IO (J (Cov (JV sx)) MX -> J (CollTraj x z u p n deg) MX -> J (CovTraj sx n) MX) mkComputeCovariances c2d computeSens qc' = do propOneCovFun <- toMXFun "propogate one covariance" (propOneCov c2d) - let computeCovs :: J (CollTrajCov sx ocp n deg) MX -> J (CovTraj sx n) MX - computeCovs collTrajCov = cat covTraj + let computeCovs :: J (Cov (JV sx)) MX -> J (CollTraj x z u p n deg) MX -> J (CovTraj sx n) MX + computeCovs p0 collTraj = cat covTraj where - CollTrajCov p0 collTraj = split collTrajCov - sensitivities = computeSens collTraj covTraj = @@ -384,7 +376,7 @@ -> IO (J (JV shr) MX -> J (JV p) MX -> J (JV x) MX -> J (Cov (JV sx)) MX -> J (JV shr) MX) mkRobustifyFunction project robustifyPathC = do proj <- toSXFun "errorSpaceProjection" $ - \(JacIn x0 x1) -> JacOut (sxCatJV (project (sxSplitJV x1) (sxSplitJV x0))) (cat JNone) + \(JacIn x0 x1) -> JacOut (catJV' (project (splitJV' x1) (splitJV' x0))) (cat JNone) let _ = proj :: SXFun (JacIn (JV sx) (J (JV x))) (JacOut (JV x) (J JNone)) @@ -401,9 +393,9 @@ (J (JV x)) (M.M (JV x) (JV sx)) - let rpc (JacIn xe parm) = JacOut (sxCatJV lol) (cat JNone) + let rpc (JacIn xe parm) = JacOut (catJV' lol) (cat JNone) where - lol = robustifyPathC (sxSplitJV x) (sxSplitJV e) (sxSplitJV parm) + lol = robustifyPathC (splitJV' x) (splitJV' e) (splitJV' parm) JTuple x e = split xe robustH <- toSXFun "robust constraint" rpc let _ = robustH :: SXFun diff --git a/src/Dyno/DirectCollocation/Types.hs b/src/Dyno/DirectCollocation/Types.hs --- a/src/Dyno/DirectCollocation/Types.hs +++ b/src/Dyno/DirectCollocation/Types.hs @@ -1,8 +1,9 @@ {-# OPTIONS_GHC -Wall #-} -{-# Language ScopedTypeVariables #-} -{-# Language DeriveGeneric #-} -{-# Language FlexibleContexts #-} -{-# Language PolyKinds #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE DeriveFunctor #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE PolyKinds #-} module Dyno.DirectCollocation.Types ( CollTraj(..) @@ -12,8 +13,6 @@ , CollStageConstraints(..) , CollOcpConstraints' , CollOcpConstraints(..) - , CollTrajCov(..) - , CollOcpCovConstraints(..) , fillCollTraj , fillCollTraj' , fmapCollTraj @@ -24,21 +23,36 @@ , fmapCollPointJ , fillCollConstraints , getXzus + , getXzus' + , getXzus'' + , getXzus''' + , fromXzus + -- * for callbacks + , Quadratures(..) + , StageOutputs(..) + , StageOutputs' + -- * robust + , CollTrajCov(..) + , CollOcpCovConstraints(..) ) where -import GHC.Generics ( Generic ) +import GHC.Generics ( Generic, Generic1 ) -import qualified Data.Foldable as F import Linear.V ( Dim(..) ) import Data.Vector ( Vector ) +import Data.Serialize ( Serialize ) +import Accessors ( Lookup ) + import Dyno.Ocp import Dyno.View.Viewable ( Viewable ) import Dyno.View.View ( View(..), J, jfill ) import Dyno.View.JVec ( JVec(..), jreplicate ) import Dyno.View.Cov ( Cov ) import Dyno.View.JV ( JV, splitJV, catJV ) -import Dyno.Vectorize ( Vectorize(..), Id ) +import Dyno.Vectorize ( Vectorize(..), Id(..) ) +import Dyno.TypeVecs ( Vec ) +import qualified Dyno.TypeVecs as TV -- | CollTraj using type families to compress type parameters @@ -110,21 +124,74 @@ , View sh, Vectorize shr, View sc ) => View (CollOcpCovConstraints ocp n deg sh shr sc) +-- todo(greg): unit test to ensure this is the inverse of getXzus' +fromXzus :: forall x z u p n deg a + . (Vectorize x, Vectorize z, Vectorize u, Vectorize p, Dim n, Dim deg) + => a -> p a -> Vec n (x a, Vec deg (x a, z a, u a)) -> x a + -> CollTraj x z u p n deg (Vector a) +fromXzus t p xzus xf = CollTraj (catJV (Id t)) (catJV p) (cat (JVec traj)) (catJV xf) + where + traj :: Vec n (J (CollStage (JV x) (JV z) (JV u) deg) (Vector a)) + traj = fmap (cat . toCollStage) xzus + toCollStage :: (x a, Vec deg (x a, z a, u a)) -> CollStage (JV x) (JV z) (JV u) deg (Vector a) + toCollStage (x0, xzus') = CollStage (catJV x0) (cat (JVec (fmap toCollPoint xzus'))) + + toCollPoint :: (x a, z a, u a) -> J (CollPoint (JV x) (JV z) (JV u)) (Vector a) + toCollPoint (x,z,u) = cat $ CollPoint (catJV x) (catJV z) (catJV u) + getXzus :: (Vectorize x, Vectorize z, Vectorize u, Dim n, Dim deg) - => CollTraj x z u p n deg (Vector a) -> ([[x a]], [[z a]], [[u a]]) -getXzus (CollTraj _ _ stages xf) = (xs ++ [[splitJV xf]], zs, us) + => CollTraj x z u p n deg (Vector a) + -> (Vec n (Vec deg (x a, z a, u a))) +getXzus traj = fmap snd $ fst $ getXzus' traj + +getXzus' :: + (Vectorize x, Vectorize z, Vectorize u, Dim n, Dim deg) + => CollTraj x z u p n deg (Vector a) + -> (Vec n (x a, Vec deg (x a, z a, u a)), x a) +getXzus' (CollTraj _ _ stages xf) = + (fmap (getXzusFromStage . split) (unJVec (split stages)), splitJV xf) + +getXzus'' :: + forall x z u p n deg a + . (Vectorize x, Vectorize z, Vectorize u, Dim n, Dim deg) + => CollTraj x z u p n deg (Vector a) + -> ( Vec n (Vec deg (x a)) + , Vec n (Vec deg (z a)) + , Vec n (Vec deg (u a)) + ) +getXzus'' traj = (fmap snd xs, zs, us) where - (xs, zs, us) = unzip3 $ map (getXzus' . split) (F.toList (unJVec (split stages))) + ((xs,_),zs,us) = getXzus''' traj -getXzus' :: (Vectorize x, Vectorize z, Vectorize u, Dim deg) - => CollStage (JV x) (JV z) (JV u) deg (Vector a) -> ([x a], [z a], [u a]) -getXzus' (CollStage x0 xzus) = (splitJV x0 : xs, zs, us) +getXzus''' :: + forall x z u p n deg a + . (Vectorize x, Vectorize z, Vectorize u, Dim n, Dim deg) + => CollTraj x z u p n deg (Vector a) + -> ( ( Vec n (x a, Vec deg (x a)) + , x a + ) + , Vec n (Vec deg (z a)) + , Vec n (Vec deg (u a)) + ) +getXzus''' traj = ((xs, xf), zs, us) where - (xs, zs, us) = unzip3 $ map (f . split) (F.toList (unJVec (split xzus))) + (xzus, xf) = getXzus' traj + (xs, zs, us) = TV.tvunzip3 $ fmap f xzus + where + f (x0, xzus') = ((x0,xs'), zs', us') + where + (xs',zs',us') = TV.tvunzip3 xzus' + +getXzusFromStage :: (Vectorize x, Vectorize z, Vectorize u, Dim deg) + => CollStage (JV x) (JV z) (JV u) deg (Vector a) + -> (x a, Vec deg (x a, z a, u a)) +getXzusFromStage (CollStage x0 xzus) = (splitJV x0, fmap (f . split) (unJVec (split xzus))) + where f (CollPoint x z u) = (splitJV x, splitJV z, splitJV u) + fillCollConstraints :: forall x r c h n deg a . ( Vectorize x, Vectorize r, Vectorize c, Vectorize h @@ -274,3 +341,35 @@ -> CollPoint x1 z1 u1 a -> CollPoint x2 z2 u2 b fmapCollPointJ fx fz fu (CollPoint x z u) = CollPoint (fx x) (fz z) (fu u) + +-- | for callbacks +data Quadratures q qo a = + Quadratures + { qLagrange :: a + , qUser :: q a + , qOutputs :: qo a + } deriving (Functor, Generic, Generic1) +instance (Vectorize q, Vectorize qo) => Vectorize (Quadratures q qo) +instance (Lookup a, Lookup (q a), Lookup (qo a)) => Lookup (Quadratures q qo a) +instance (Serialize a, Serialize (q a), Serialize (qo a)) => Serialize (Quadratures q qo a) + +-- | for callbacks +data StageOutputs x o h q qo po deg a = + StageOutputs + { soVec :: Vec deg ( J (JV o) (Vector a) + , J (JV x) (Vector a) + , J (JV h) (Vector a) + , J (JV po) (Vector a) + , Quadratures q qo a -- qs + , Quadratures q qo a -- qdots + ) + , soXNext :: J (JV x) (Vector a) + , soQNext :: Quadratures q qo a + } deriving Generic + +type StageOutputs' ocp deg = StageOutputs (X ocp) (O ocp) (H ocp) (Q ocp) (QO ocp) (PO ocp) deg + +instance ( Serialize a, Serialize (q a), Serialize (qo a) + , Vectorize x, Vectorize o, Vectorize h, Vectorize po + , Dim deg + ) => (Serialize (StageOutputs x o h q qo po deg a)) diff --git a/src/Dyno/FormatTime.hs b/src/Dyno/FormatTime.hs new file mode 100644 --- /dev/null +++ b/src/Dyno/FormatTime.hs @@ -0,0 +1,62 @@ +{-# OPTIONS_GHC -Wall #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE DeriveFunctor #-} +{-# LANGUAGE DeriveGeneric #-} + +module Dyno.FormatTime + ( formatSeconds + ) where + +import Text.Printf ( printf ) + +-- | format seconds in a more human readable way +-- +-- >>> formatSeconds 59 +-- "59.0 seconds" +-- +-- >>> formatSeconds 59.999 +-- "1 minute, 0.0 seconds" +-- +-- >>> formatSeconds 1000 +-- "16 minutes, 40.0 seconds" +-- +-- >>> formatSeconds 1019.99 +-- "17 minutes, 0.0 seconds" +-- +-- >>> formatSeconds 3599.9 +-- "59 minutes, 59.9 seconds" +-- +-- >>> formatSeconds 3599.99 +-- "1 hour, 0 minutes, 0.0 seconds" +-- +-- >>> formatSeconds 3600 +-- "1 hour, 0 minutes, 0.0 seconds" +-- +-- >>> formatSeconds 123456 +-- "34 hours, 17 minutes, 36.0 seconds" +formatSeconds :: Double -> String +formatSeconds seconds' = formatHMCS hours minutes (round $ 10 * seconds) + where + hours :: Int + hours = floor $ seconds' / (60 * 60) + + minutes :: Int + minutes = floor $ seconds' / 60 - (fromIntegral hours) * 60 + + seconds :: Double + seconds = seconds' - (fromIntegral hours)*60*60 - (fromIntegral minutes)*60 + + formatHMCS :: Int -> Int -> Int -> String + -- handle printf rounding up + formatHMCS h m 600 = formatHMCS h (m + 1) 0 + formatHMCS h 60 cs = formatHMCS (h+1) 0 cs + -- format bumped numbers + formatHMCS 0 0 cs = printf "%.1f seconds" (fromCS cs) + formatHMCS 0 1 cs = printf "1 minute, %.1f seconds" (fromCS cs) + formatHMCS 0 m cs = printf "%d minutes, %.1f seconds" m (fromCS cs) + formatHMCS 1 1 cs = printf "1 hour, 1 minute, %.1f seconds" (fromCS cs) + formatHMCS 1 m cs = printf "1 hour, %d minutes, %.1f seconds" m (fromCS cs) + formatHMCS h m cs = printf "%d hours, %d minutes, %.1f seconds" h m (fromCS cs) + + fromCS :: Int -> Double + fromCS = (0.1 *) . fromIntegral diff --git a/src/Dyno/Integrate.hs b/src/Dyno/Integrate.hs new file mode 100644 --- /dev/null +++ b/src/Dyno/Integrate.hs @@ -0,0 +1,42 @@ +{-# OPTIONS_GHC -Wall #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} + +module Dyno.Integrate + ( InitialTime(..) + , TimeStep(..) + , rk45 + ) where + +import qualified Data.Vector as V +import qualified Data.Vector.Storable as SV +import qualified Numeric.GSL.ODE as ODE +import qualified Numeric.LinearAlgebra.Data as D + +import Dyno.Vectorize ( Vectorize(..), devectorize ) + +newtype InitialTime = InitialTime {unInitialTime :: Double} + deriving (Num, Fractional, Floating, Ord, Eq, Show) +newtype TimeStep = TimeStep {unTimeStep :: Double} + deriving (Num, Fractional, Floating, Ord, Eq, Show) + +rk45 :: Vectorize x + => (Double -> x Double -> x Double) + -> InitialTime -> TimeStep -> x Double -> x Double +rk45 f (InitialTime t0) (TimeStep h) x0 = devectorize $ sv $ last sol + where + vs :: V.Vector Double -> SV.Vector Double + vs = SV.fromList . V.toList + sv :: SV.Vector Double -> V.Vector Double + sv = V.fromList . SV.toList + + sol = D.toRows $ + ODE.odeSolveV + ODE.RKf45 +-- ODE.RK8pd +-- ODE.MSAdams --todo(benchmark) + h 1e-8 1e-6 f' + (vs (vectorize x0)) + (SV.fromList [0, h]) + f' :: Double -> SV.Vector Double -> SV.Vector Double + f' t x = vs $ vectorize $ f (t + t0) (devectorize (sv x)) diff --git a/src/Dyno/LagrangePolynomials.lhs b/src/Dyno/LagrangePolynomials.lhs --- a/src/Dyno/LagrangePolynomials.lhs +++ b/src/Dyno/LagrangePolynomials.lhs @@ -123,7 +123,7 @@ import Dyno.TypeVecs -interpolate :: (Additive f, Fractional a) => Vec deg a -> Vec deg (f a) -> a -> f a +interpolate :: (Additive f, Dim deg, Fractional a) => Vec deg a -> Vec deg (f a) -> a -> f a interpolate taus0 xs0 tau1 = sumV [x ^* (lagrangeXis taus0' tau1 k) | (k,x) <- zip [0..] xs0'] where taus0' = F.toList taus0 diff --git a/src/Dyno/MultipleShooting.hs b/src/Dyno/MultipleShooting.hs --- a/src/Dyno/MultipleShooting.hs +++ b/src/Dyno/MultipleShooting.hs @@ -1,7 +1,7 @@ {-# OPTIONS_GHC -Wall #-} -{-# Language ScopedTypeVariables #-} -{-# Language DeriveGeneric #-} -{-# Language PolyKinds #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE PolyKinds #-} module Dyno.MultipleShooting ( MsOcp(..) @@ -78,7 +78,7 @@ k1 = f x0 u p t k2 = f (x0 ^+^ h/2 *^ k1) u p (t+h/2) k3 = f (x0 ^+^ h/2 *^ k2) u p (t+h/2) - k4 = f (x0 ^+^ h *^ k2) u p (t+h) + k4 = f (x0 ^+^ h *^ k3) u p (t+h) simulate :: (Floating a, Additive x) => Int -> Ode x u p a -> x a -> u a -> p a -> a -> a -> x a simulate n ode x0' u p t h = xf diff --git a/src/Dyno/Nlp.hs b/src/Dyno/Nlp.hs --- a/src/Dyno/Nlp.hs +++ b/src/Dyno/Nlp.hs @@ -1,7 +1,7 @@ {-# OPTIONS_GHC -Wall #-} -{-# Language FlexibleInstances #-} -{-# Language DeriveFunctor #-} -{-# Language DeriveGeneric #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE DeriveFunctor #-} +{-# LANGUAGE DeriveGeneric #-} module Dyno.Nlp ( Bounds diff --git a/src/Dyno/NlpScaling.hs b/src/Dyno/NlpScaling.hs --- a/src/Dyno/NlpScaling.hs +++ b/src/Dyno/NlpScaling.hs @@ -1,5 +1,5 @@ {-# OPTIONS_GHC -Wall #-} -{-# Language ScopedTypeVariables #-} +{-# LANGUAGE ScopedTypeVariables #-} module Dyno.NlpScaling ( ScaleFuns(..) diff --git a/src/Dyno/NlpSolver.hs b/src/Dyno/NlpSolver.hs --- a/src/Dyno/NlpSolver.hs +++ b/src/Dyno/NlpSolver.hs @@ -1,13 +1,15 @@ {-# OPTIONS_GHC -Wall #-} -{-# Language ScopedTypeVariables #-} -{-# Language PackageImports #-} -{-# Language KindSignatures #-} -{-# Language GeneralizedNewtypeDeriving #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE PackageImports #-} +{-# LANGUAGE KindSignatures #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} module Dyno.NlpSolver ( NlpSolver - , SXElement , runNlpSolver + , RunNlpOptions(..) + , runNlpSolverWith + , defaultRunnerOptions -- * solve , solve , solve' @@ -58,11 +60,13 @@ , generateAndCompile ) where +import Text.Printf ( printf ) +import Data.Time.Clock ( getCurrentTime, diffUTCTime ) import Data.Proxy ( Proxy(..) ) import System.Process ( callProcess, showCommandForUser ) import Control.Exception ( AsyncException( UserInterrupt ), try ) import Control.Concurrent ( forkIO, newEmptyMVar, takeMVar, putMVar ) -import Control.Applicative ( Applicative(..) ) +import qualified Control.Applicative as A import Control.Monad ( when, void ) import "mtl" Control.Monad.Reader ( MonadIO(..), MonadReader(..), ReaderT(..) ) import Data.Maybe ( fromMaybe ) @@ -88,7 +92,7 @@ import Dyno.View.Unsafe.View ( unJ, mkJ ) import Dyno.View.Unsafe.M ( mkM ) -import Dyno.SXElement ( SXElement ) +import Dyno.FormatTime ( formatSeconds ) import Dyno.Vectorize ( Id(..) ) import Dyno.View.JV ( JV ) import Dyno.View.View ( View(..), J, fmapJ, d2v, v2d, jfill ) @@ -98,11 +102,19 @@ import Dyno.View.Viewable ( Viewable ) import Dyno.Nlp ( NlpOut(..), KKT(..) ) import Dyno.NlpScaling ( ScaleFuns(..), scaledFG, mkScaleFuns ) -import Dyno.Solvers ( Solver(..) ) +import Dyno.Solvers ( Solver(..), getSolverInternal ) +import Dyno.SolverInternal ( SolverInternal(..) ) type VD a = J a (Vector Double) type VMD a = J a (Vector (Maybe Double)) +timeIt :: IO a -> IO (a, Double) +timeIt action = do + t0 <- getCurrentTime + ret <- action + t1 <- getCurrentTime + return (ret, realToFrac (diffUTCTime t1 t0)) + getStat :: String -> NlpSolver x p g C.GenericType getStat name = do nlpState <- ask @@ -111,7 +123,7 @@ setInput :: View xg => (ScaleFuns x g DMatrix -> (J xg DMatrix -> J xg DMatrix)) - -> (NlpState x g -> Int) + -> (NlpState x p g -> Int) -> String -> J xg (V.Vector Double) -> NlpSolver x p g () @@ -151,7 +163,10 @@ setUbg = setInput gToGBar isNg "ubg" . toUb setP :: View p => VD p -> NlpSolver x p g () -setP = setInput (const id) isNp "p" +setP p = do + nlpState <- ask + isSetParam nlpState p + setInput (const id) isNp "p" p setLamX0 :: View x => VD x -> NlpSolver x p g () setLamX0 = setInput lamXToLamXBar isNx "lam_x0" @@ -465,7 +480,7 @@ return nlpOut -data NlpState (x :: * -> *) (g :: * -> *) = +data NlpState (x :: * -> *) (p :: * -> *) (g :: * -> *) = NlpState { isNx :: Int , isNg :: Int @@ -474,13 +489,14 @@ , isInterrupt :: IO () , isSuccessCodes :: [String] , isScale :: ScaleFuns x g DMatrix + , isSetParam :: J p (Vector Double) -> NlpSolver x p g () } newtype NlpSolver (x :: * -> *) (p :: * -> *) (g :: * -> *) a = - NlpSolver (ReaderT (NlpState x g) IO a) + NlpSolver (ReaderT (NlpState x p g) IO a) deriving ( Functor - , Applicative + , A.Applicative , Monad - , MonadReader (NlpState x g) + , MonadReader (NlpState x p g) , MonadIO ) @@ -495,6 +511,17 @@ callProcess cmd args externalFunction ("./"++name++".so") +data RunNlpOptions = + RunNlpOptions + { verbose :: Bool + } + +defaultRunnerOptions :: RunNlpOptions +defaultRunnerOptions = + RunNlpOptions + { verbose = False + } + runNlpSolver :: forall x p g a s . (View x, View p, View g, Symbolic s) @@ -503,10 +530,24 @@ -> Maybe (J x (Vector Double)) -> Maybe (J g (Vector Double)) -> Maybe Double - -> Maybe (J x (Vector Double) -> IO Bool) + -> Maybe (J x (Vector Double) -> J p (Vector Double) -> IO Bool) -> NlpSolver x p g a -> IO a -runNlpSolver solverStuff nlpFun scaleX scaleG scaleF callback' (NlpSolver nlpMonad) = do +runNlpSolver = runNlpSolverWith defaultRunnerOptions + +runNlpSolverWith :: + forall x p g a s . + (View x, View p, View g, Symbolic s) + => RunNlpOptions + -> Solver + -> (J x s -> J p s -> (J (JV Id) s, J g s)) + -> Maybe (J x (Vector Double)) + -> Maybe (J g (Vector Double)) + -> Maybe Double + -> Maybe (J x (Vector Double) -> J p (Vector Double) -> IO Bool) + -> NlpSolver x p g a + -> IO a +runNlpSolverWith runnerOptions solverStuff nlpFun scaleX scaleG scaleF callback' (NlpSolver nlpMonad) = do inputsX <- sym "x" inputsP <- sym "p" @@ -522,12 +563,21 @@ inputScheme <- mkScheme SCHEME_NLPInput [("x", inputsXMat), ("p", inputsPMat)] outputScheme <- mkScheme SCHEME_NLPOutput [("f", objMat), ("g", gMat)] - nlp <- mkFunction "nlp" inputScheme outputScheme --- Op.setOption nlp "verbose" True + + when (verbose runnerOptions) $ do + putStrLn "************** initializing dynobud runNlpSolver ******************" + putStrLn "making nlp..." + (nlp, nlpTime) <- timeIt $ mkFunction "nlp" inputScheme outputScheme + when (verbose runnerOptions) $ printf "made nlp in %s\n" (formatSeconds nlpTime) mapM_ (\(l,Op.Opt o) -> Op.setOption nlp l o) (functionOptions solverStuff) - soInit nlp + when (verbose runnerOptions) $ putStrLn "init nlp..." + (_, nlpInitTime) <- timeIt $ soInit nlp + when (verbose runnerOptions) $ printf "nlp initialized in %s\n" (formatSeconds nlpInitTime) - functionCall solverStuff nlp + when (verbose runnerOptions) $ putStrLn "function call..." + -- in case the user wants to do something (like codegen?) + (_, functionCallTime) <- timeIt $ functionCall solverStuff nlp + when (verbose runnerOptions) $ printf "function called in %s\n" (formatSeconds functionCallTime) -- let eval 0 = error "finished" -- eval k = do @@ -544,20 +594,24 @@ -- jac_sparsity <- C.function_jacSparsity nlp 0 1 True False -- C.sparsity_spyMatlab jac_sparsity "jac_sparsity_reorder.m" - - solver <- C.nlpSolver__0 (solverName solverStuff) nlp + when (verbose runnerOptions) $ putStrLn "create solver..." + (solver, solverCreateTime) <- timeIt $ C.nlpSolver__0 (solverName (getSolverInternal solverStuff)) nlp + when (verbose runnerOptions) $ printf "created solver in %s\n" (formatSeconds solverCreateTime) -- add callback if user provides it intref <- newIORef False + paramRef <- newIORef (jfill 0) let cb function' = do callbackRet <- case callback' of Nothing -> return True Just callback -> do xval <- fmap (d2v . xbarToX scale . mkJ . CM.densify) $ C.ioInterfaceFunction_output__2 function' 0 - callback xval + pval <- readIORef paramRef + callback xval pval interrupt <- readIORef intref - return $ if callbackRet && not interrupt then 0 else fromIntegral (solverInterruptCode solverStuff) + return $ if callbackRet && not interrupt then 0 + else fromIntegral (solverInterruptCode (getSolverInternal solverStuff)) casadiCallback <- makeCallback cb >>= C.genericType__0 Op.setOption solver "iteration_callback" casadiCallback -- grad_f <- gradient nlp 0 0 @@ -582,8 +636,11 @@ -- Op.setOption solver "jac_g" jac_g' -- set all the user options - mapM_ (\(l,Op.Opt o) -> Op.setOption solver l o) (defaultOptions solverStuff ++ options solverStuff) - soInit solver + mapM_ (\(l,Op.Opt o) -> Op.setOption solver l o) (defaultSolverOptions (getSolverInternal solverStuff) + ++ options solverStuff) + when (verbose runnerOptions) $ putStrLn "initialize solver..." + (_, solverInitTime) <- timeIt $ soInit solver + when (verbose runnerOptions) $ printf "solver initialized in %s\n" (formatSeconds solverInitTime) let proxy :: J f b -> Proxy f proxy = const Proxy @@ -593,7 +650,12 @@ , isNg = size (proxy g) , isSolver = solver , isInterrupt = writeIORef intref True - , isSuccessCodes = successCodes solverStuff + , isSuccessCodes = successCodes (getSolverInternal solverStuff) , isScale = scale + , isSetParam = liftIO . writeIORef paramRef } - liftIO $ runReaderT nlpMonad nlpState + when (verbose runnerOptions) $ putStrLn "run NLP monad..." + (ret, retTime) <- timeIt $ liftIO $ runReaderT nlpMonad nlpState + when (verbose runnerOptions) $ printf "ran NLP monad in %s\n" (formatSeconds retTime) + return ret + diff --git a/src/Dyno/NlpUtils.hs b/src/Dyno/NlpUtils.hs --- a/src/Dyno/NlpUtils.hs +++ b/src/Dyno/NlpUtils.hs @@ -1,18 +1,20 @@ {-# OPTIONS_GHC -Wall #-} -{-# Language ScopedTypeVariables #-} -{-# Language RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE RankNTypes #-} module Dyno.NlpUtils ( HomotopyParams(..) , solveNlpHomotopy + , solveNlpHomotopyWith , solveNlp + , solveNlpWith , solveNlpV , setNlpInputs , runNlp + , runNlpWith ) where -import Control.Applicative ( Applicative(..) ) -import Data.Maybe ( fromMaybe ) +import qualified Control.Applicative as A import qualified Data.Traversable as T import Control.Monad ( when, void ) import Data.Vector ( Vector ) @@ -27,7 +29,7 @@ import Dyno.Vectorize ( Vectorize(..), Id(..) ) import Dyno.View.JV ( JV, catJV, catJV', splitJV, splitJV' ) -import Dyno.View.View ( View(..), J, JNone(..), JTuple(..), jfill, unzipJ, fmapJ ) +import Dyno.View.View ( View(..), J, JNone(..), unzipJ ) import Dyno.View.Symbolic ( Symbolic ) import Dyno.Nlp ( Nlp(..), NlpOut(..), Bounds ) import Dyno.Solvers ( Solver ) @@ -39,7 +41,7 @@ fmap f (StateL k) = StateL $ \ s -> do (s', v) <- k s return (s', f v) -instance Monad m => Applicative (StateL m s) where +instance Monad m => A.Applicative (StateL m s) where pure x = StateL (\s -> return (s, x)) StateL kf <*> StateL kv = StateL $ \ s -> do (s', f) <- kf s @@ -65,40 +67,49 @@ (View x, View p, View g, T.Traversable t, Symbolic a) => Double -> HomotopyParams -> Solver - -> Maybe (J p (Vector Double)) - -> Nlp x p g a -> t (J p (Vector Double)) -> Maybe (J (JTuple x p) (Vector Double) -> IO Bool) + -> Nlp x p g a -> t (J p (Vector Double)) + -> Maybe (J x (Vector Double) -> J p (Vector Double) -> IO Bool) -> Maybe (J x (Vector Double) -> J p (Vector Double) -> Double -> IO ()) - -> IO (t (NlpOut (JTuple x p) g (Vector Double))) -solveNlpHomotopy userStep hp - solverStuff pscale nlp pFs callback callbackP = do + -> IO (t (NlpOut x g (Vector Double))) +solveNlpHomotopy = solveNlpHomotopyWith defaultRunnerOptions + +-- | solve a homotopy nlp +solveNlpHomotopyWith :: + forall x p g t a . + (View x, View p, View g, T.Traversable t, Symbolic a) + => RunNlpOptions + -> Double -> HomotopyParams + -> Solver + -> Nlp x p g a -> t (J p (Vector Double)) + -> Maybe (J x (Vector Double) -> J p (Vector Double) -> IO Bool) + -> Maybe (J x (Vector Double) -> J p (Vector Double) -> Double -> IO ()) + -> IO (t (NlpOut x g (Vector Double))) +solveNlpHomotopyWith options userStep hp + solverStuff nlp pFs callback callbackP = do when ((reduction hp) >= 1) $ error $ "homotopy reduction factor " ++ show (reduction hp) ++ " >= 1" when ((increase hp) <= 1) $ error $ "homotopy increase factor " ++ show (increase hp) ++ " <= 1" - let fg :: J (JTuple x p) a -> J JNone a -> (J (JV Id) a, J g a) - fg xp _ = nlpFG nlp x p - where - JTuple x p = split xp + let fg :: J x a -> J p a -> (J (JV Id) a, J g a) + fg x p = nlpFG nlp x p - xpscale :: Maybe (J (JTuple x p) (Vector Double)) - xpscale = case (nlpScaleX nlp, pscale) of - (Nothing, Nothing) -> Nothing - (xs, ps) -> Just $ cat $ JTuple (fromMaybe (jfill 1) xs) (fromMaybe (jfill 1) ps) - runNlpSolver solverStuff fg xpscale (nlpScaleG nlp) (nlpScaleF nlp) callback $ do + runNlpSolverWith options solverStuff fg (nlpScaleX nlp) (nlpScaleG nlp) (nlpScaleF nlp) callback $ do let (lbx,ubx) = unzipJ (nlpBX nlp) (lbg,ubg) = unzipJ (nlpBG nlp) p0 = nlpP nlp - setBnds p' = do - setLbx $ cat (JTuple lbx (fmapJ Just p')) - setUbx $ cat (JTuple ubx (fmapJ Just p')) + _ <- case callback of + Nothing -> return True + Just cb -> liftIO $ cb (nlpX0 nlp) (nlpP nlp) -- initial solve - setX0 $ cat $ JTuple (nlpX0 nlp) (nlpP nlp) - setP $ cat JNone - setBnds p0 + setX0 $ nlpX0 nlp + setP $ nlpP nlp + setLbx lbx + setUbx ubx setLbg lbg setUbg ubg + -- todo(greg): clean up redundancy? case nlpLamX0 nlp of - Just lam -> setLamX0 $ cat (JTuple lam (jfill 0)) + Just lam -> setLamX0 lam Nothing -> return () case nlpLamG0 nlp of Just lam -> setLamG0 lam @@ -115,15 +126,15 @@ let runCallback alphaTrial = case callbackP of Nothing -> return () Just cbp -> do - xp <- getX - let JTuple x p = split xp + x <- getX + p <- getP liftIO $ void (cbp x p alphaTrial) let solveOneStage :: (Int, Double, J p (Vector Double)) -> J p (Vector Double) - -> NlpSolver (JTuple x p) JNone g - ((Int, Double, J p (Vector Double)), NlpOut (JTuple x p) g (Vector Double)) + -> NlpSolver x p g + ((Int, Double, J p (Vector Double)), NlpOut x g (Vector Double)) solveOneStage (stage, step0, p0') pF' = do ((msg, ret'), stepF) <- tryStep 0 0 step0 ret <- case msg of @@ -131,15 +142,15 @@ Right _ -> return ret' return ((stage + 1, stepF, pF'), ret) where - setAlpha :: Double -> NlpSolver (JTuple x p) JNone g () + setAlpha :: Double -> NlpSolver x p g () setAlpha alpha = do let p0'' = unJ p0' let p = mkJ $ V.zipWith (+) p0'' (V.map (alpha*) (V.zipWith (-) (unJ pF') p0'')) - setBnds p + setP p tryStep :: Int -> Double -> Double - -> NlpSolver (JTuple x p) JNone g - ((Either String String, NlpOut (JTuple x p) g (Vector Double)), Double) + -> NlpSolver x p g + ((Either String String, NlpOut x g (Vector Double)), Double) tryStep majorIter alpha0 step | step < 1e-12 = do _no <- getNlpOut error "step size too small" @@ -218,8 +229,10 @@ -- :: Maybe (J (JV g) (V.Vector Double)) } - callback :: Maybe (J (JV x) (Vector Double) -> IO Bool) - callback = fmap (. splitJV) cb + callback :: Maybe (J (JV x) (Vector Double) -> J JNone (Vector Double) -> IO Bool) + callback = case cb of + Nothing -> Nothing + Just cb' -> Just $ \x _ -> cb' (splitJV x) (r0, r1) <- solveNlp solverStuff nlp callback return $ case r0 of @@ -241,12 +254,22 @@ solveNlp :: (View x, View p, View g, Symbolic a) => Solver - -> Nlp x p g a -> Maybe (J x (Vector Double) -> IO Bool) + -> Nlp x p g a -> Maybe (J x (Vector Double) -> J p (Vector Double) -> IO Bool) -> IO (Either String String, NlpOut x g (Vector Double)) solveNlp solverStuff nlp callback = runNlp solverStuff nlp callback solve' +-- | convenience function to solve a pure Nlp +solveNlpWith :: + (View x, View p, View g, Symbolic a) + => RunNlpOptions + -> Solver + -> Nlp x p g a -> Maybe (J x (Vector Double) -> J p (Vector Double) -> IO Bool) + -> IO (Either String String, NlpOut x g (Vector Double)) +solveNlpWith opts solverStuff nlp callback = + runNlpWith opts solverStuff nlp callback solve' + -- | set all inputs setNlpInputs :: (View x, View p, View g, Symbolic a) => Nlp x p g a -> NlpSolver x p g () setNlpInputs nlp = do @@ -271,10 +294,20 @@ runNlp :: (View x, View p, View g, Symbolic a) => Solver - -> Nlp x p g a -> Maybe (J x (Vector Double) -> IO Bool) + -> Nlp x p g a -> Maybe (J x (Vector Double) -> J p (Vector Double) -> IO Bool) -> NlpSolver x p g b -> IO b -runNlp solverStuff nlp callback runMe = - runNlpSolver solverStuff (nlpFG nlp) (nlpScaleX nlp) (nlpScaleG nlp) (nlpScaleF nlp) callback $ do +runNlp = runNlpWith defaultRunnerOptions + +-- | set all inputs, handle scaling, and let the user run a NlpMonad +runNlpWith :: + (View x, View p, View g, Symbolic a) + => RunNlpOptions + -> Solver + -> Nlp x p g a -> Maybe (J x (Vector Double) -> J p (Vector Double) -> IO Bool) + -> NlpSolver x p g b + -> IO b +runNlpWith options solverStuff nlp callback runMe = + runNlpSolverWith options solverStuff (nlpFG nlp) (nlpScaleX nlp) (nlpScaleG nlp) (nlpScaleF nlp) callback $ do setNlpInputs nlp runMe diff --git a/src/Dyno/Ocp.hs b/src/Dyno/Ocp.hs --- a/src/Dyno/Ocp.hs +++ b/src/Dyno/Ocp.hs @@ -1,10 +1,14 @@ {-# OPTIONS_GHC -Wall #-} -{-# Language TypeFamilies #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE FlexibleContexts #-} module Dyno.Ocp ( OcpPhase(..) + , OcpPhaseInputs(..) , OcpPhaseWithCov(..) , OcpPhase' + , OcpPhaseInputs' , X , Z , U @@ -14,22 +18,27 @@ , C , H , Q + , QO + , PO + , FP ) where +import GHC.Generics ( Generic ) + +import Data.Serialize ( Serialize ) import Data.Vector ( Vector ) import Dyno.View.JV ( JV ) import Dyno.View.View ( J ) import Dyno.View.Cov ( Cov ) import Dyno.Nlp ( Bounds ) -import Dyno.SXElement ( SXElement ) ---import Dyno.Vectorize +import Dyno.Vectorize ( Id ) import Casadi.SX ( SX ) import Casadi.DMatrix ( DMatrix ) type Sx a = J a SX -type Sxe = SXElement +type Sxe = J (JV Id) SX -- | differential state type family X a :: * -> * @@ -49,14 +58,21 @@ type family H a :: * -> * -- | quadrature state type family Q a :: * -> * +-- | quadrature output states (for plotting only) +type family QO a :: * -> * +-- | plot outputs +type family PO a :: * -> * +-- | fixed (hardcoded) parameters +type family FP a :: * -> * -- | OcpPhase using type families to compress type parameters -type OcpPhase' ocp = OcpPhase (X ocp) (Z ocp) (U ocp) (P ocp) (R ocp) (O ocp) (C ocp) (H ocp) (Q ocp) +type OcpPhase' ocp = OcpPhase (X ocp) (Z ocp) (U ocp) (P ocp) (R ocp) (O ocp) (C ocp) (H ocp) (Q ocp) (QO ocp) (PO ocp) (FP ocp) +type OcpPhaseInputs' ocp = OcpPhaseInputs (X ocp) (Z ocp) (U ocp) (P ocp) (C ocp) (H ocp) (FP ocp) -- | One stage of an optimal control problem, solvable as a stand-alone optimal control problem. -- --- > minimize Jm(x(T),T) + integrate( Jl(x(t),z(t),u(t),p,t), {t,0,T} ) +-- > minimize Jm(x(T),T) + integrate( Jl(x(t),z(t),u(t),p,p',t), {t,0,T} ) -- > x(.), z(.), u(.), p, T -- > -- > subject to: @@ -70,49 +86,47 @@ -- -- nonlinear path constraints -- --- > hlb <= h(x(t), z(t), u(t), p, t) <= hub +-- > hlb <= h(x(t), z(t), u(t), p, p', t) <= hub -- -- dynamics constraints: -- --- > f(x'(t), x(t), z(t), u(t), p, t) == 0 +-- > f(x'(t), x(t), z(t), u(t), p, p', t) == 0 -- -- boundary conditions: -- --- > c(x(0), x(T), q(T), p) == 0 +-- > c(x(0), x(T), q(T), p, p') == 0 -- -- perhaps this should be: -- -- > c(x(0), 0, x(T), T) == 0 -data OcpPhase x z u p r o c h q = +-- +-- +-- The OcpPhase data type has all the symbolics necessary to set up a problem. +-- The OcpPhaseInputs data type provides bounds on all parameters. +-- It is split up this way because setting up a problem takes considerable overhead, +-- so solving many problem with different OcpPhaseInputs can save time. +data OcpPhase x z u p r o c h q qo po fp = OcpPhase - { -- | the Mayer term @Jm(T, x(0), x(T), q(T), p)@ - ocpMayer :: Sxe -> x Sxe -> x Sxe -> q Sxe -> p Sxe -> Sxe - -- | the Lagrange term @Jl(x(t),z(t),u(t),p,o,t,T)@ - , ocpLagrange :: x Sxe -> z Sxe -> u Sxe -> p Sxe -> o Sxe -> Sxe -> Sxe -> Sxe - -- | derivative of quadrature state @q(x(t),z(t),u(t),p,o,t,T)@ - , ocpQuadratures :: x Sxe -> z Sxe -> u Sxe -> p Sxe -> o Sxe -> Sxe -> Sxe -> q Sxe + { -- | the Mayer term @Jm(T, x(0), x(T), q(T), p, p')@ + ocpMayer :: Sxe -> x Sxe -> x Sxe -> q Sxe -> p Sxe -> fp Sxe -> Sxe + -- | the Lagrange term @Jl(x(t),z(t),u(t),p,p',o,t,T)@ + , ocpLagrange :: x Sxe -> z Sxe -> u Sxe -> p Sxe -> fp Sxe -> o Sxe -> Sxe -> Sxe -> Sxe + -- | derivative of quadrature state @q(x(t),z(t),u(t),p,o,p',t,T)@ + , ocpQuadratures :: x Sxe -> z Sxe -> u Sxe -> p Sxe -> fp Sxe -> o Sxe -> Sxe -> Sxe -> q Sxe + -- | same as ocpQuadratures, but only used for plotting + , ocpQuadratureOutputs :: x Sxe -> z Sxe -> u Sxe -> p Sxe -> fp Sxe -> o Sxe -> Sxe -> Sxe -> qo Sxe -- | fully implicit differential-algebraic equation of the form: -- - -- > f(x'(t), x(t), z(t), u(t), p, t) == 0 - , ocpDae :: x Sxe -> x Sxe -> z Sxe -> u Sxe -> p Sxe -> Sxe -> (r Sxe, o Sxe) - -- | the boundary conditions @clb <= c(x(0), x(T), q(T), T) <= cub@ - , ocpBc :: x Sxe -> x Sxe -> q Sxe -> p Sxe -> Sxe -> c Sxe - -- | the path constraints @h(x(t), z(t), u(t), p, t)@ - , ocpPathC :: x Sxe -> z Sxe -> u Sxe -> p Sxe -> o Sxe -> Sxe -> h Sxe - -- | the boundary condition bounds @clb <= c(x(0), x(T)) <= cub@ - , ocpBcBnds :: c Bounds - -- | the path constraint bounds @(hlb, hub)@ - , ocpPathCBnds :: h Bounds - -- | differential state bounds @(xlb, xub)@ - , ocpXbnd :: x Bounds - -- | algebraic variable bounds @(zlb, zub)@ - , ocpZbnd :: z Bounds - -- | control bounds @(ulb, uub)@ - , ocpUbnd :: u Bounds - -- | parameter bounds @(plb, pub)@ - , ocpPbnd :: p Bounds - -- | time bounds @(Tlb, Tub)@ - , ocpTbnd :: Bounds + -- > f(x'(t), x(t), z(t), u(t), p, p', t) == 0 + , ocpDae :: x Sxe -> x Sxe -> z Sxe -> u Sxe -> p Sxe -> fp Sxe -> Sxe -> (r Sxe, o Sxe) + -- | the boundary conditions @clb <= c(x(0), x(T), q(T), p, p', T) <= cub@ + , ocpBc :: x Sxe -> x Sxe -> q Sxe -> p Sxe -> fp Sxe -> Sxe -> c Sxe + -- | the path constraints @hbl <= h(x(t), z(t), u(t), p, p', o, t) <= hbu@ + , ocpPathC :: x Sxe -> z Sxe -> u Sxe -> p Sxe -> fp Sxe -> o Sxe -> Sxe -> h Sxe + -- | things you might want to plot, like total energy - integral(power) + -- + -- > po((x(0), x(T)), x(t), z(t), u(t), p, o(t), q(t), qo(t), fp, t, T) + , ocpPlotOutputs :: (x Sxe, x Sxe) -> x Sxe -> z Sxe -> u Sxe -> p Sxe -> o Sxe -> q Sxe -> qo Sxe -> fp Sxe -> Sxe -> Sxe -> po Sxe -- | scaling , ocpObjScale :: Maybe Double , ocpTScale :: Maybe Double @@ -125,6 +139,37 @@ , ocpPathCScale :: Maybe (h Double) } + +-- | Inputs to an OcpPhase problem, used to solve several different problems with one OcpPhase. +data OcpPhaseInputs x z u p c h fp = + OcpPhaseInputs + { -- | the boundary condition bounds @clb <= c(x(0), x(T), q(T), p, p', T) <= cub@ + ocpBcBnds :: c Bounds + -- | the path constraint bounds @hbl <= h(x(t), z(t), u(t), p, p', o, t) <= hbu@ + , ocpPathCBnds :: h Bounds + -- | differential state bounds @xlb <= x(t) <= xub@ + , ocpXbnd :: x Bounds + -- | algebraic variable bounds @zlb <= z(t) <= zub@ + , ocpZbnd :: z Bounds + -- | control bounds @ulb <= u(t) <= uub@ + , ocpUbnd :: u Bounds + -- | parameter bounds @plb <= p <= pub@ + , ocpPbnd :: p Bounds + -- | time bounds @Tlb <= T <= Tub@ + , ocpTbnd :: Bounds + -- | fixed parameters (not optimization variables) + , ocpFixedP :: fp Double + } deriving ( Generic ) + +instance ( Serialize (x Bounds) + , Serialize (z Bounds) + , Serialize (u Bounds) + , Serialize (p Bounds) + , Serialize (c Bounds) + , Serialize (h Bounds) + , Serialize (fp Double) + ) => Serialize (OcpPhaseInputs x z u p c h fp) + data OcpPhaseWithCov ocp sx sz sw sr sh shr sc = OcpPhaseWithCov { -- | the Mayer term @Jm(T, x(0), x(T), P(0), P(t))@ @@ -151,7 +196,7 @@ , ocpCovSbc :: Sx (Cov (JV sx)) -> Sx (Cov (JV sx)) -> Sx sc , ocpCovSbcBnds :: J sc (Vector Bounds) -- | the covariance path constraints @h(s)@, only applied to first n Ss - , ocpCovSh :: X ocp SXElement -> Sx (Cov (JV sx)) -> Sx sh + , ocpCovSh :: X ocp Sxe -> Sx (Cov (JV sx)) -> Sx sh , ocpCovShBnds :: J sh (Vector Bounds) -- | scaling , ocpCovSScale :: Maybe (J (Cov (JV sx)) (Vector Double)) diff --git a/src/Dyno/OcpHomotopy.hs b/src/Dyno/OcpHomotopy.hs --- a/src/Dyno/OcpHomotopy.hs +++ b/src/Dyno/OcpHomotopy.hs @@ -1,9 +1,10 @@ {-# OPTIONS_GHC -Wall #-} -{-# Language ScopedTypeVariables #-} -{-# Language PolyKinds #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE PolyKinds #-} module Dyno.OcpHomotopy ( runOcpHomotopy + , runOcpHomotopyWith ) where import Data.Vector ( Vector ) @@ -12,173 +13,107 @@ import Casadi.MX ( MX ) import Dyno.Ocp -import Dyno.Vectorize ( Vectorize, Id(..), None(..) ) -import Dyno.View.View ( View(..), J, JTuple(..), JNone(..) ) -import Dyno.View.JV ( JV, catJV, catJV' ) -import Dyno.View.Viewable ( Viewable ) +import Dyno.Vectorize ( Vectorize ) +import Dyno.View.View ( J ) +import Dyno.View.JV ( JV, catJV ) import Dyno.TypeVecs ( Dim ) import Dyno.Solvers ( Solver ) import Dyno.Nlp ( Nlp(..), NlpOut(..) ) -import Dyno.NlpUtils ( HomotopyParams(..), solveNlp, solveNlpHomotopy ) +import Dyno.NlpSolver ( RunNlpOptions, defaultRunnerOptions ) +import Dyno.NlpUtils ( HomotopyParams(..), solveNlpWith, solveNlpHomotopyWith ) import Dyno.DirectCollocation.Types ( CollTraj(..), CollOcpConstraints ) import Dyno.DirectCollocation.Formulate ( CollProblem(..), makeCollProblem ) import Dyno.DirectCollocation.Quadratures ( QuadratureRoots ) -tupleToCollTraj :: - forall x z u p n deg a - . ( Dim deg, Dim n, Viewable a, Vectorize x, Vectorize z, Vectorize u, Vectorize p ) - => JTuple (CollTraj x z u None n deg) (JV p) a - -> J (CollTraj x z u p n deg) a -tupleToCollTraj (JTuple x0 p) = cat x1 - where - x1 :: CollTraj x z u p n deg a - x1 = CollTraj tf p stages0 xf - - CollTraj tf _ stages0 xf = split x0 - -collTrajToTuple :: - forall x z u p n deg a . - ( Viewable a - , Vectorize x, Vectorize z, Vectorize u, Vectorize p - , Dim deg, Dim n ) - => J (CollTraj x z u p n deg) a - -> JTuple (CollTraj x z u None n deg) (JV p) a -collTrajToTuple x0 = JTuple (cat x1) pfp - where - x1 :: CollTraj x z u None n deg a - x1 = CollTraj tf (catJV' None) stages0 xf - - CollTraj tf pfp stages0 xf = split x0 - - -convertNlp :: - forall x z u p r c h n deg a - . ( Viewable a - , Vectorize x, Vectorize z, Vectorize u, Vectorize p - , Dim deg, Dim n - ) - => Nlp (CollTraj x z u p n deg) JNone (CollOcpConstraints x r c h n deg) a - -> Nlp (CollTraj x z u None n deg) (JV p) (CollOcpConstraints x r c h n deg) a -convertNlp nlp0 = nlp - where - nlp = Nlp { nlpX0 = x0 - , nlpBG = nlpBG nlp0 - , nlpP = fp0 - , nlpFG = fg - , nlpBX = bx - , nlpLamX0 = fmap ((\(JTuple ret _) -> ret) . collTrajToTuple) (nlpLamX0 nlp0) - , nlpLamG0 = nlpLamG0 nlp0 - , nlpScaleF = nlpScaleF nlp0 - , nlpScaleX = fmap ((\(JTuple ret _) -> ret) . collTrajToTuple) (nlpScaleX nlp0) - , nlpScaleG = nlpScaleG nlp0 - } - fg :: J (CollTraj x z u None n deg) a -> J (JV p) a - -> (J (JV Id) a, J (CollOcpConstraints x r c h n deg) a) - fg x fp = (f,g') - where - (f,g') = nlpFG nlp0 x' (cat JNone) - x' :: J (CollTraj x z u p n deg) a - x' = tupleToCollTraj (JTuple x fp) - - JTuple x0 fp0 = collTrajToTuple (nlpX0 nlp0) - JTuple bx _ = collTrajToTuple (nlpBX nlp0) - - - -runOcpHomotopy :: - forall x z u p r o c h q n deg t +runOcpHomotopyWith :: + forall x z u p r o c h q qo po fp n deg t . ( Dim n, Dim deg , Vectorize x, Vectorize z, Vectorize u, Vectorize p - , Vectorize r, Vectorize o, Vectorize c, Vectorize h, Vectorize q + , Vectorize r, Vectorize o, Vectorize c, Vectorize h + , Vectorize q, Vectorize po, Vectorize qo + , Vectorize fp , T.Traversable t ) - => Double -> HomotopyParams - -> OcpPhase x z u p r o c h q + => RunNlpOptions + -> Double -> HomotopyParams + -> OcpPhase x z u p r o c h q qo po fp + -> OcpPhaseInputs x z u p c h fp -> J (CollTraj x z u p n deg) (Vector Double) - -> QuadratureRoots -> Bool -> Bool -> Solver -> Solver -> p Double -> t (p Double) - -> (CollProblem x z u p r o c h q n deg - -> IO ([String] -> J (CollTraj x z u p n deg) (Vector Double) -> IO Bool) + -> QuadratureRoots -> Bool -> Bool -> Solver -> Solver + -> t (fp Double) + -> (CollProblem x z u p r o c h q qo po fp n deg + -> IO ([String] -> J (CollTraj x z u p n deg) (Vector Double) -> J (JV fp) (Vector Double) -> IO Bool) ) -> IO (t (NlpOut (CollTraj x z u p n deg) (CollOcpConstraints x r c h n deg) (Vector Double))) -runOcpHomotopy step0 homotopyParams ocpHomotopy guess roots useStartupCallback useHomotopyCallback - startupSolver homotopySolver param0 nominalParams makeCallback = do - cp0 <- makeCollProblem roots ocpHomotopy guess - callbackHeh <- makeCallback cp0 - let nlp0 = cpNlp cp0 +runOcpHomotopyWith opts step0 homotopyParams ocpHomotopy ocpHomotopyInputs guess roots + useStartupCallback useHomotopyCallback + startupSolver homotopySolver nominalParams makeCallback = do + cp0 <- makeCollProblem roots ocpHomotopy ocpHomotopyInputs guess + callback <- makeCallback cp0 let nlpHomotopy :: Nlp - (CollTraj x z u None n deg) - (JV p) + (CollTraj x z u p n deg) + (JV fp) (CollOcpConstraints x r c h n deg) MX - nlpHomotopy = (convertNlp nlp0) {nlpP = catJV param0} - - let callback :: [String] - -> J (CollTraj x z u p n deg) (Vector Double) - -> IO Bool - callback moarMsgs traj = callbackHeh moarMsgs traj + nlpHomotopy = cpNlp cp0 - scb = if useStartupCallback then Just (callback ["homotopy startup solve"]) else Nothing + let scb = if useStartupCallback + then Just (callback ["homotopy startup solve"]) + else Nothing putStrLn "running startup solver..." - (msg0,opt0') <- solveNlp startupSolver nlp0 scb + (msg0,opt0') <- solveNlpWith opts startupSolver nlpHomotopy scb opt0 <- case msg0 of Left msg' -> error msg' Right _ -> return opt0' - let homoGuessX :: J (CollTraj x z u None n deg) (Vector Double) - JTuple homoGuessX _ = collTrajToTuple $ xOpt opt0 - JTuple homoGuessLX _ = collTrajToTuple $ lambdaXOpt opt0 + let homoGuessX :: J (CollTraj x z u p n deg) (Vector Double) + homoGuessX = xOpt opt0 + homoGuessLX = lambdaXOpt opt0 homoGuessLG :: J (CollOcpConstraints x r c h n deg) (Vector Double) homoGuessLG = lambdaGOpt opt0 - pFinals :: t (J (JV p) (Vector Double)) + pFinals :: t (J (JV fp) (Vector Double)) pFinals = fmap catJV nominalParams - homoCallback :: J (JTuple (CollTraj x z u None n deg) (JV p)) (Vector Double) + homoCallback :: J (CollTraj x z u p n deg) (Vector Double) -> J (JV fp) (Vector Double) -> IO Bool - homoCallback traj0 = callback [ "homotopy stepping" - ] (tupleToCollTraj (split traj0)) + homoCallback traj0 fp = + callback [ "homotopy stepping" + ] traj0 fp putStrLn "\ninitial solve done, starting homotopy steps" let hcb = if useHomotopyCallback then Just homoCallback else Nothing - pscale :: Maybe (J (JV p) (Vector Double)) - pscale = fmap catJV (ocpPScale ocpHomotopy) - opt1s <- solveNlpHomotopy step0 homotopyParams - homotopySolver - pscale - (nlpHomotopy { nlpX0 = homoGuessX - , nlpLamX0 = Just homoGuessLX - , nlpLamG0 = Just homoGuessLG - }) - pFinals - hcb Nothing - - let f :: NlpOut (JTuple (CollTraj x z u None n deg) (JV p)) - (CollOcpConstraints x r c h n deg) - (Vector Double) - -> NlpOut (CollTraj x z u p n deg) - (CollOcpConstraints x r c h n deg) - (Vector Double) - f nlpOut = - NlpOut - { fOpt = fOpt nlpOut - , xOpt = g (xOpt nlpOut) - , gOpt = gOpt nlpOut - , lambdaXOpt = g (lambdaXOpt nlpOut) - , lambdaGOpt = lambdaGOpt nlpOut - } - where - g :: J (JTuple (CollTraj x z u None n deg) (JV p)) (Vector Double) - -> J (CollTraj x z u p n deg) (Vector Double) - g = tupleToCollTraj . split - - ret :: t (NlpOut - (CollTraj x z u p n deg) - (CollOcpConstraints x r c h n deg) - (Vector Double)) - ret = fmap f opt1s + solveNlpHomotopyWith opts step0 homotopyParams + homotopySolver + (nlpHomotopy { nlpX0 = homoGuessX + , nlpLamX0 = Just homoGuessLX + , nlpLamG0 = Just homoGuessLG + }) + pFinals + hcb Nothing - return ret +runOcpHomotopy :: + forall x z u p r o c h q qo po fp n deg t + . ( Dim n, Dim deg + , Vectorize x, Vectorize z, Vectorize u, Vectorize p + , Vectorize r, Vectorize o, Vectorize c, Vectorize h + , Vectorize q, Vectorize po, Vectorize qo + , Vectorize fp + , T.Traversable t ) + => Double -> HomotopyParams + -> OcpPhase x z u p r o c h q qo po fp + -> OcpPhaseInputs x z u p c h fp + -> J (CollTraj x z u p n deg) (Vector Double) + -> QuadratureRoots -> Bool -> Bool -> Solver -> Solver + -> t (fp Double) + -> (CollProblem x z u p r o c h q qo po fp n deg + -> IO ([String] -> J (CollTraj x z u p n deg) (Vector Double) -> J (JV fp) (Vector Double) -> IO Bool) + ) + -> IO (t (NlpOut (CollTraj x z u p n deg) + (CollOcpConstraints x r c h n deg) + (Vector Double))) +runOcpHomotopy = runOcpHomotopyWith defaultRunnerOptions diff --git a/src/Dyno/SXElement.hs b/src/Dyno/SXElement.hs deleted file mode 100644 --- a/src/Dyno/SXElement.hs +++ /dev/null @@ -1,64 +0,0 @@ -{-# OPTIONS_GHC -Wall #-} -{-# Language GeneralizedNewtypeDeriving #-} - -module Dyno.SXElement - ( SXElement(..) - , sxSplitJV - , sxCatJV - -- todo: remove this completely after NlpMonad/OcpMonad are done with it - , sxElementSym - -- todo: remove the next two exports after NlpMonad/OcpMonad are done with it - , sxElementToSX - , sxToSXElement - ) where - -import Linear.Conjugate ( Conjugate(..) ) - -import Casadi.SX ( SX, ssym ) -import qualified Casadi.CMatrix as CM -import Casadi.Overloading ( Fmod, ArcTan2, SymOrd ) - -import Dyno.View.Unsafe.View ( mkJ, unJ ) - -import Dyno.View.JV ( JV, splitJV', catJV' ) -import Dyno.View.View ( J ) -import Dyno.Vectorize ( Vectorize, Id ) - -newtype SXElement = - SXElement SX - deriving ( Num, Fractional, Floating - , Fmod, ArcTan2, SymOrd - , Show, Eq, Conjugate - ) - --- todo: take this out after NlpMonad/OcpMonad are done with it -sxElementSym :: String -> IO SXElement -sxElementSym = fmap SXElement . ssym - - -sxToSXElement :: SX -> SXElement -sxToSXElement x - | (1,1) == sizes = SXElement x - | otherwise = error $ "sxToSXElement: got non-scalar of size " ++ show sizes - where - sizes = (CM.size1 x, CM.size2 x) - -sxElementToSX :: SXElement -> SX -sxElementToSX (SXElement x) - | (1,1) == sizes = x - | otherwise = error $ "sxElementToSX: got non-scalar of size " ++ show sizes - where - sizes = (CM.size1 x, CM.size2 x) - - -sxSplitJV :: Vectorize f => J (JV f) SX -> f SXElement -sxSplitJV v = fmap f (splitJV' v) - where - f :: J (JV Id) SX -> SXElement - f = sxToSXElement . unJ - -sxCatJV :: Vectorize f => f SXElement -> J (JV f) SX -sxCatJV v = catJV' (fmap f v) - where - f :: SXElement -> J (JV Id) SX - f x = mkJ (sxElementToSX x) diff --git a/src/Dyno/SimpleOcp.hs b/src/Dyno/SimpleOcp.hs new file mode 100644 --- /dev/null +++ b/src/Dyno/SimpleOcp.hs @@ -0,0 +1,133 @@ +{-# OPTIONS_GHC -Wall #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE DeriveFunctor #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} + +module Dyno.SimpleOcp + ( SimpleOcp(..) + , S + , solveOcp + ) where + +import GHC.Generics ( Generic, Generic1 ) + +import qualified Data.Foldable as F +import Data.Proxy +import Data.Vector ( Vector ) + +import Accessors ( Lookup ) +import Casadi.SX ( SX ) + +import Dyno.TypeVecs +import Dyno.Ocp +import Dyno.Solvers +import Dyno.NlpUtils +import Dyno.Nlp +import Dyno.DirectCollocation.Formulate +import Dyno.DirectCollocation.Types +import Dyno.DirectCollocation.Quadratures +import Dyno.Vectorize ( Vectorize(..), Tuple(..), Id, None(..), fill, vzipWith ) +import Dyno.View.View -- ( View(..) ) +import Dyno.View.JV +import Dyno.View.JVec + +-- | scalar symbolic type +newtype S = S {unS :: J (JV Id) SX} deriving (Num, Fractional, Floating) + +data SimpleOcp x u = + SimpleOcp + { ode :: x S -> u S -> x S + , objective :: x S -> u S -> S + , xBounds :: x (Double, Double) + , uBounds :: u (Double, Double) + , xInitial :: x Double + , xFinal :: x Double + , endTime :: Double + , initialGuess :: Double -> x Double + } + +vminus :: (Vectorize f, Num a) => f a -> f a -> f a +vminus = vzipWith (-) + +dot :: (Vectorize f, Num a) => f a -> f a -> a +dot x y = F.sum $ vectorize $ vzipWith (*) x y + +toOcp :: (Vectorize x, Vectorize u) + => SimpleOcp x u + -> OcpPhase (Tuple x u) None u None (Tuple x u) None (SimpleBc x) None None None None None +toOcp simple = + OcpPhase + { ocpMayer = \_ _ _ _ _ _ -> 0 + , ocpLagrange = \(Tuple x u) _ u' _ _ _ _ _ -> 1e-9 * (u' `dot` u') + unS (objective simple (fmap S x) (fmap S u)) + , ocpQuadratures = \_ _ _ _ _ _ _ _ -> None + , ocpQuadratureOutputs = \_ _ _ _ _ _ _ _ -> None + , ocpDae = \(Tuple xd ud) (Tuple x u) _ u' _ _ _ -> + let r = Tuple (xd `vminus` x') (ud `vminus` u') + x' = fmap unS $ ode simple (fmap S x) (fmap S u) + in (r, None) + , ocpBc = \(Tuple x0 _) (Tuple xf _) _ _ _ _ ->SimpleBc x0 xf + , ocpPathC = \_ _ _ _ _ _ _ -> None + , ocpPlotOutputs = \_ _ _ _ _ _ _ _ _ _ _ -> None + , ocpObjScale = Nothing + , ocpTScale = Nothing + , ocpXScale = Nothing + , ocpZScale = Nothing + , ocpUScale = Nothing + , ocpPScale = Nothing + , ocpResidualScale = Nothing + , ocpBcScale = Nothing + , ocpPathCScale = Nothing + } + +data SimpleBc x a = SimpleBc (x a) (x a) deriving (Functor, Generic, Generic1) +instance Vectorize x => Vectorize (SimpleBc x) +instance Lookup (x a) => Lookup (SimpleBc x a) + +toOcpInputs :: (Vectorize x, Vectorize u) => SimpleOcp x u -> OcpPhaseInputs (Tuple x u) None u None (SimpleBc x) None None +toOcpInputs simple = + OcpPhaseInputs + { ocpBcBnds = SimpleBc + (fmap (\x -> (Just x, Just x)) (xInitial simple)) + (fmap (\x -> (Just x, Just x)) (xFinal simple)) + , ocpPathCBnds = None + , ocpXbnd = fmap toBounds $ Tuple (xBounds simple) (uBounds simple) + , ocpUbnd = fill (Nothing, Nothing) + , ocpZbnd = None + , ocpPbnd = None + , ocpTbnd = (Just (endTime simple), Just (endTime simple)) + , ocpFixedP = None + } + where + toBounds (lb,ub) = (Just lb, Just ub) + +solveOcp :: (Vectorize x, Vectorize u) => SimpleOcp x u -> IO (Either String [(x Double, u Double)]) +solveOcp simple = reifyDim deg $ reifyDim n $ solveOcp' simple + where + n = 50 + deg = 2 + +solver :: Solver +solver = ipoptSolver + +solveOcp' :: + forall x u n deg + . (Vectorize x, Vectorize u, Dim deg, Dim n) + => SimpleOcp x u -> Proxy n -> Proxy deg -> IO (Either String [(x Double, u Double)]) +solveOcp' simple _ _ = do + let ocp = toOcp simple + ocpInputs = toOcpInputs simple + tf = endTime simple + roots = Legendre + guess :: CollTraj (Tuple x u) None u None n deg (Vector Double) + guess = makeGuess roots tf (\t -> Tuple (initialGuess simple t) (fill 0)) (const None) (const (fill 0)) None + cp <- makeCollProblem roots ocp ocpInputs (cat guess) + let _ = cp :: CollProblem (Tuple x u) None u None (Tuple x u) None (SimpleBc x) None None None None None n deg + (emsg, opt) <- solveNlp solver (cpNlp cp) Nothing + case emsg of + Left msg -> return (Left msg) + Right _ -> do + let CollTraj _ _ stages' xf' = split (xOpt opt) + xs = map ((\(CollStage x _) -> splitJV x) . split) $ F.toList $ unJVec (split stages') + return $ Right $ map (\(Tuple x u) -> (x, u)) (xs ++ [splitJV xf']) diff --git a/src/Dyno/SolverInternal.hs b/src/Dyno/SolverInternal.hs new file mode 100644 --- /dev/null +++ b/src/Dyno/SolverInternal.hs @@ -0,0 +1,15 @@ +{-# OPTIONS_GHC -Wall #-} + +module Dyno.SolverInternal + ( SolverInternal(..) + ) where + +import Casadi.Option ( Opt(..) ) + +data SolverInternal = + SolverInternal + { solverName :: String + , defaultSolverOptions :: [(String,Opt)] + , solverInterruptCode :: Int + , successCodes :: [String] + } diff --git a/src/Dyno/Solvers.hs b/src/Dyno/Solvers.hs --- a/src/Dyno/Solvers.hs +++ b/src/Dyno/Solvers.hs @@ -1,82 +1,126 @@ {-# OPTIONS_GHC -Wall #-} -module Dyno.Solvers ( Solver(..) +module Dyno.Solvers ( Solver(options, functionOptions, functionCall) , Opt(..) , ipoptSolver, snoptSolver, worhpSolver + , getSolverInternal ) where import Casadi.Core.Classes.Function ( Function ) import Casadi.Option ( Opt(..) ) +import Dyno.SolverInternal ( SolverInternal(..) ) + data Solver = Solver - { solverName :: String - , defaultOptions :: [(String,Opt)] - , options :: [(String,Opt)] - , solverInterruptCode :: Int - , successCodes :: [String] + { options :: [(String,Opt)] , functionOptions :: [(String, Opt)] , functionCall :: Function -> IO () + , solverInternal :: SolverInternal } +-- | get the read-only part +getSolverInternal :: Solver -> SolverInternal +getSolverInternal = solverInternal + snoptSolver :: Solver snoptSolver = Solver - { solverName = "snopt" - , defaultOptions = [ -- ("_iprint", Opt (0::Int)) --- , ("_isumm", Opt (6::Int)) --- , ("_scale_option", Opt (0::Int)) --- , ("_major_iteration_limit", Opt (3 :: Int)) --- , ("_minor_iteration_limit", Opt (2000 :: Int)) --- , ("_verify_level", Opt (2 :: Int)) --- , ("_optimality_tolerance", Opt (1e-1 :: Double)) --- , ("_feasibility_tolerance", Opt (1e-1 :: Double)) --- , ("detect_linear", Opt False) --- , ("monitor", Opt (V.fromList ["setup_nlp"]) ) --- , ("_start", Opt "Warm") - ] - , options = [] - , solverInterruptCode = -2 - , successCodes = ["1"] + { options = [] , functionOptions = [] , functionCall = const (return ()) + , solverInternal = + SolverInternal + { solverName = "snopt" + , defaultSolverOptions = + [ -- ("_iprint", Opt (0::Int)) +-- , ("_isumm", Opt (6::Int)) +-- , ("_scale_option", Opt (0::Int)) +-- , ("_major_iteration_limit", Opt (3 :: Int)) +-- , ("_minor_iteration_limit", Opt (2000 :: Int)) +-- , ("_verify_level", Opt (2 :: Int)) +-- , ("_optimality_tolerance", Opt (1e-1 :: Double)) +-- , ("_feasibility_tolerance", Opt (1e-1 :: Double)) +-- , ("detect_linear", Opt False) +-- , ("monitor", Opt (V.fromList ["setup_nlp"]) ) +-- , ("_start", Opt "Warm") + ] + , solverInterruptCode = -2 + , successCodes = ["1"] + } } ipoptSolver :: Solver ipoptSolver = Solver - { solverName = "ipopt" - , defaultOptions = [ ("max_iter", Opt (3000 :: Int)) - , ("tol", Opt (1e-9 :: Double)) --- , ("hessian_approximation", Opt "limited-memory") --- , ("expand", Opt True) --- , ("linear_solver", Opt "ma27") --- , ("linear_solver", Opt "ma57") --- , ("linear_solver", Opt "ma86") --- , ("linear_solver", Opt "ma97") --- , ("fixed_variable_treatment", Opt "make_constraint") -- causes segfaults? --- , ("fixed_variable_treatment", Opt "make_parameter") - ] - , options = [] - , solverInterruptCode = 1 - , successCodes = ["Solve_Succeeded", "Solved_To_Acceptable_Level"] + { options = [] , functionOptions = [] , functionCall = const (return ()) + , solverInternal = + SolverInternal + { solverName = "ipopt" + , defaultSolverOptions = + [ ("max_iter", Opt (3000 :: Int)) + , ("tol", Opt (1e-9 :: Double)) +-- , ("hessian_approximation", Opt "limited-memory") +-- , ("expand", Opt True) +-- , ("linear_solver", Opt "ma27") +-- , ("linear_solver", Opt "ma57") +-- , ("linear_solver", Opt "ma86") +-- , ("linear_solver", Opt "ma97") +-- , ("fixed_variable_treatment", Opt "make_constraint") -- causes segfaults? +-- , ("fixed_variable_treatment", Opt "make_parameter") + ] + , solverInterruptCode = 1 + , successCodes = ["Solve_Succeeded", "Solved_To_Acceptable_Level"] + } } worhpSolver :: Solver worhpSolver = Solver - { solverName = "worhp" - , defaultOptions = [] - , options = [] - , solverInterruptCode = 1 - , successCodes = [ "OptimalSolution" - , "LowPassFilterOptimal" - ] + { options = [] , functionOptions = [] , functionCall = const (return ()) + , solverInternal = + SolverInternal + { solverName = "worhp" + , defaultSolverOptions = [] + , solverInterruptCode = 1 + , successCodes = [ "OptimalSolution" + , "LowPassFilterOptimal" + ] + } } + + + + +--_sqpSolver :: Solver +--_sqpSolver = +-- Solver +-- { solverName = "sqpmethod" +-- , defaultOptions = [ ("qp_solver", Opt "nlp") +-- , ("qp_solver_options" +-- , Opt [ ( "nlp_solver", Opt "ipopt") +-- , ( "nlp_solver_options" +-- , Opt [ ("tol", Opt (1e-12 :: Double)) +-- , ("linear_solver", Opt "ma86") +-- , ("ma86_order", Opt "metis") +-- , ("print_level", Opt (0 :: Int)) +-- , ("print_time", Opt False) +-- ] +-- ) +-- ] +-- ) +-- ] +-- , options = [] +-- , solverInterruptCode = 1 +-- , successCodes = [""] +-- , functionOptions = [] +-- , functionCall = const (return ()) +-- } +-- diff --git a/src/Dyno/TypeVecs.hs b/src/Dyno/TypeVecs.hs --- a/src/Dyno/TypeVecs.hs +++ b/src/Dyno/TypeVecs.hs @@ -8,12 +8,12 @@ {-# LANGUAGE GeneralizedNewtypeDeriving #-} {-# LANGUAGE UndecidableInstances #-} {-# LANGUAGE PolyKinds #-} -- so that "Vec (n :: Nat) a" works +{-# LANGUAGE InstanceSigs #-} module Dyno.TypeVecs ( Vec , Succ , unVec - , mkVec , mkVec' , tvlength , (|>) @@ -46,8 +46,6 @@ import GHC.Generics ( Generic, Generic1 ) import Control.Applicative -import Data.Foldable ( Foldable ) -import Data.Traversable ( Traversable ) import qualified Data.Traversable as T import qualified Data.Vector as V import Data.Vector.Binary () -- instances @@ -59,6 +57,7 @@ import Data.Proxy import Data.Reflection as R import Data.Distributive ( Distributive(..) ) +import Prelude -- BBP workaround import Accessors ( Lookup(..), AccessorTree(..) ) @@ -66,25 +65,37 @@ -- length-indexed vectors using phantom types newtype Vec (n :: k) a = MkVec (V.Vector a) - deriving (Eq, Ord, Functor, Traversable, Foldable, Generic, Generic1) + deriving (Functor, Generic, Generic1) instance (Dim n, B.Binary a) => B.Binary (Vec n a) where put = B.put . unVec - get = fmap mkVec B.get + get = do + x <- B.get + case devectorize' x of + Right y -> return y + Left msg -> fail msg instance (Dim n, S.Serialize a) => S.Serialize (Vec n a) where put = S.put . unVec - get = fmap mkVec S.get + get = do + x <- S.get + case devectorize' x of + Right y -> return y + Left msg -> fail msg instance (Lookup a, Dim n) => Lookup (Vec n a) where - toAccessorTree vec f = Data ("Vec " ++ show n, "Vec " ++ show n) $ map child (take n [0..]) + toAccessorTree vec get set = Data ("Vec " ++ show n, "Vec " ++ show n) $ map child (take n [0..]) where n = reflectDim (Proxy :: Proxy n) - child k = ("v" ++ show k, toAccessorTree (getK vec) (getK . f)) + child k = ("v" ++ show k, toAccessorTree (getK vec) (getK . get) setK) where + setK vk new = set (devectorize (v V.// [(k,vk)])) new + where + MkVec v = get new + getK :: Vec n a -> a getK (MkVec v) = v V.! k instance Dim n => Distributive (Vec n) where - distribute f = mkVec $ V.generate (reflectDim (Proxy :: Proxy n)) + distribute f = devectorize $ V.generate (reflectDim (Proxy :: Proxy n)) $ \i -> fmap (\v -> V.unsafeIndex (vectorize v) i) f {-# INLINE distribute #-} @@ -107,12 +118,18 @@ MkVec xs ^-^ MkVec ys = MkVec (V.zipWith (-) xs ys) instance Dim n => Vectorize (Vec n) where + fill = pure vectorize = unVec - devectorize = mkVec - empty = pure () + devectorize' :: V.Vector a -> Either String (Vec n a) + devectorize' x + | n == n' = Right (MkVec x) + | otherwise = Left $ "mkVec: length mismatch, " ++ show (n,n') + where + n = reflectDim (Proxy :: Proxy n) + n' = V.length x tvtranspose :: (Dim n, Dim m) => Vec n (Vec m a) -> Vec m (Vec n a) -tvtranspose vec = mkVec $ fmap mkVec $ T.sequence (unVec (fmap unVec vec)) +tvtranspose = T.sequenceA infixr 5 <| infixl 5 |> @@ -130,70 +147,61 @@ n = reflectDim (Proxy :: Proxy n) n' = V.length x --- todo: put these in unsafe module -mkVec :: forall n a . Dim n => V.Vector a -> Vec n a -mkVec x - | n == n' = MkVec x - | otherwise = error $ "mkVec: length mismatch, " ++ show (n,n') - where - n = reflectDim (Proxy :: Proxy n) - n' = V.length x - mkVec' :: Dim n => [a] -> Vec n a -mkVec' = mkVec . V.fromList +mkVec' = devectorize . V.fromList tvlength :: forall n a. Dim n => Vec n a -> Int tvlength = const $ reflectDim (Proxy :: Proxy n) tvzip :: Dim n => Vec n a -> Vec n b -> Vec n (a,b) -tvzip x y = mkVec (V.zip (unVec x) (unVec y)) +tvzip x y = devectorize (V.zip (unVec x) (unVec y)) tvzip3 :: Dim n => Vec n a -> Vec n b -> Vec n c -> Vec n (a,b,c) -tvzip3 x y z = mkVec (V.zip3 (unVec x) (unVec y) (unVec z)) +tvzip3 x y z = devectorize (V.zip3 (unVec x) (unVec y) (unVec z)) tvzip4 :: Dim n => Vec n a -> Vec n b -> Vec n c -> Vec n d -> Vec n (a,b,c,d) -tvzip4 x y z w = mkVec (V.zip4 (unVec x) (unVec y) (unVec z) (unVec w)) +tvzip4 x y z w = devectorize (V.zip4 (unVec x) (unVec y) (unVec z) (unVec w)) tvzipWith :: Dim n => (a -> b -> c) -> Vec n a -> Vec n b -> Vec n c -tvzipWith f x y = mkVec (V.zipWith f (unVec x) (unVec y)) +tvzipWith f x y = devectorize (V.zipWith f (unVec x) (unVec y)) tvzipWith3 :: Dim n => (a -> b -> c -> d) -> Vec n a -> Vec n b -> Vec n c -> Vec n d -tvzipWith3 f x y z = mkVec (V.zipWith3 f (unVec x) (unVec y) (unVec z)) +tvzipWith3 f x y z = devectorize (V.zipWith3 f (unVec x) (unVec y) (unVec z)) tvzipWith4 :: Dim n => (a -> b -> c -> d -> e) -> Vec n a -> Vec n b -> Vec n c -> Vec n d -> Vec n e -tvzipWith4 f x y z u = mkVec (V.zipWith4 f (unVec x) (unVec y) (unVec z) (unVec u)) +tvzipWith4 f x y z u = devectorize (V.zipWith4 f (unVec x) (unVec y) (unVec z) (unVec u)) tvzipWith5 :: Dim n => (a -> b -> c -> d -> e -> f) -> Vec n a -> Vec n b -> Vec n c -> Vec n d -> Vec n e -> Vec n f tvzipWith5 f x0 x1 x2 x3 x4 = - mkVec (V.zipWith5 f (unVec x0) (unVec x1) (unVec x2) (unVec x3) (unVec x4)) + devectorize (V.zipWith5 f (unVec x0) (unVec x1) (unVec x2) (unVec x3) (unVec x4)) tvzipWith6 :: Dim n => (a -> b -> c -> d -> e -> f -> g) -> Vec n a -> Vec n b -> Vec n c -> Vec n d -> Vec n e -> Vec n f -> Vec n g tvzipWith6 f x0 x1 x2 x3 x4 x5 = - mkVec (V.zipWith6 f (unVec x0) (unVec x1) (unVec x2) (unVec x3) (unVec x4) (unVec x5)) + devectorize (V.zipWith6 f (unVec x0) (unVec x1) (unVec x2) (unVec x3) (unVec x4) (unVec x5)) tvunzip :: Dim n => Vec n (a,b) -> (Vec n a, Vec n b) -tvunzip v = (mkVec v1, mkVec v2) +tvunzip v = (devectorize v1, devectorize v2) where (v1,v2) = V.unzip (unVec v) tvunzip3 :: Dim n => Vec n (a,b,c) -> (Vec n a, Vec n b, Vec n c) -tvunzip3 v = (mkVec v1, mkVec v2, mkVec v3) +tvunzip3 v = (devectorize v1, devectorize v2, devectorize v3) where (v1,v2,v3) = V.unzip3 (unVec v) tvunzip4 :: Dim n => Vec n (a,b,c,d) -> (Vec n a, Vec n b, Vec n c, Vec n d) -tvunzip4 v = (mkVec v1, mkVec v2, mkVec v3, mkVec v4) +tvunzip4 v = (devectorize v1, devectorize v2, devectorize v3, devectorize v4) where (v1,v2,v3,v4) = V.unzip4 (unVec v) tvunzip5 :: Dim n => Vec n (a,b,c,d,e) -> (Vec n a, Vec n b, Vec n c, Vec n d, Vec n e) -tvunzip5 v = (mkVec v1, mkVec v2, mkVec v3, mkVec v4, mkVec v5) +tvunzip5 v = (devectorize v1, devectorize v2, devectorize v3, devectorize v4, devectorize v5) where (v1,v2,v3,v4,v5) = V.unzip5 (unVec v) @@ -207,7 +215,7 @@ tvtail :: Dim n => Vec (Succ n) a -> Vec n a tvtail x = case V.length v of 0 -> error "tvtail: empty" - _ -> mkVec $ V.tail v + _ -> devectorize $ V.tail v where v = unVec x @@ -219,10 +227,10 @@ v = unVec x tvshiftl :: Dim n => Vec n a -> a -> Vec n a -tvshiftl xs x = mkVec $ V.tail (V.snoc (unVec xs) x) +tvshiftl xs x = devectorize $ V.tail (V.snoc (unVec xs) x) tvshiftr :: Dim n => a -> Vec n a -> Vec n a -tvshiftr x xs = mkVec $ V.init (V.cons x (unVec xs)) +tvshiftr x xs = devectorize $ V.init (V.cons x (unVec xs)) instance Show a => Show (Vec n a) where showsPrec _ (MkVec v) = showV (V.toList v) @@ -250,7 +258,7 @@ {-# INLINE reifyDim #-} reifyVector :: forall a r. V.Vector a -> (forall (n :: *). Dim n => Vec n a -> r) -> r -reifyVector v f = reifyDim (V.length v) $ \(Proxy :: Proxy n) -> f (mkVec v :: Vec n a) +reifyVector v f = reifyDim (V.length v) $ \(Proxy :: Proxy n) -> f (devectorize v :: Vec n a) {-# INLINE reifyVector #-} tvlinspace :: forall n a . (Dim n, Fractional a) => a -> a -> Vec n a diff --git a/src/Dyno/Vectorize.hs b/src/Dyno/Vectorize.hs --- a/src/Dyno/Vectorize.hs +++ b/src/Dyno/Vectorize.hs @@ -12,10 +12,14 @@ {-# OPTIONS_GHC -fno-warn-orphans #-} {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE UndecidableInstances #-} +{-# LANGUAGE CPP #-} +#if defined(__GLASGOW_HASKELL__) && __GLASGOW_HASKELL__ <= 708 {-# LANGUAGE OverlappingInstances #-} +#endif module Dyno.Vectorize ( Vectorize(..) + , devectorize , None(..) , Id(..) , Tuple(..) @@ -24,48 +28,52 @@ , vzipWith , vzipWith3 , vzipWith4 - , fill , GVectorize(..) ) where import GHC.Generics -import Control.Applicative ( Applicative(..), (<$>) ) +import Control.Applicative +import Data.Either ( partitionEithers ) +import Data.Serialize ( Serialize ) import qualified Data.Vector as V -import Data.Foldable ( Foldable ) import qualified Data.Foldable as F -import Data.Traversable ( Traversable ) import qualified Data.Traversable as T import Data.Proxy ( Proxy(..) ) import qualified Linear +import Text.Printf ( printf ) +import Prelude -- BBP workaround import SpatialMath ( Euler ) import SpatialMathT ( V3T, Rot ) import Accessors ( Lookup ) + -- | a length-0 vectorizable type data None a = None - deriving (Eq, Ord, Generic, Generic1, Functor, Foldable, Traversable, Show) + deriving (Eq, Ord, Generic, Generic1, Functor, F.Foldable, T.Traversable, Show) instance Vectorize None instance Applicative None where pure = const None (<*>) = const (const None) instance Linear.Additive None where +instance Serialize (None a) -- | a length-1 vectorizable type newtype Id a = Id { unId :: a } - deriving (Eq, Ord, Generic, Generic1, Functor, Foldable, Traversable, Show) + deriving (Eq, Ord, Generic, Generic1, Functor, F.Foldable, T.Traversable, Show) instance Vectorize Id instance Applicative Id where pure = Id Id fx <*> Id x = Id (fx x) instance Linear.Additive Id where +instance Serialize a => Serialize (Id a) -- | a length-2 vectorizable type -data Tuple f g a = Tuple (f a) (g a) - deriving (Eq, Ord, Generic, Generic1, Functor, Foldable, Traversable, Show) +data Tuple f g a = Tuple { unFst :: f a, unSnd :: g a } + deriving (Eq, Ord, Generic, Generic1, Functor, F.Foldable, T.Traversable, Show) instance (Vectorize f, Vectorize g) => Vectorize (Tuple f g) instance (Applicative f, Applicative g) => Applicative (Tuple f g) where pure x = Tuple (pure x) (pure x) @@ -75,8 +83,8 @@ -- | a length-3 vectorizable type -data Triple f g h a = Triple (f a) (g a) (h a) - deriving (Eq, Ord, Generic, Generic1, Functor, Foldable, Traversable, Show) +data Triple f g h a = Triple { unFst3 :: f a, unSnd3 :: g a, unThd3 :: h a } + deriving (Eq, Ord, Generic, Generic1, Functor, F.Foldable, T.Traversable, Show) instance (Vectorize f, Vectorize g, Vectorize h) => Vectorize (Triple f g h) instance (Applicative f, Applicative g, Applicative h) => Applicative (Triple f g h) where pure x = Triple (pure x) (pure x) (pure x) @@ -104,90 +112,133 @@ instance Vectorize (V3T f) instance Vectorize (Rot f1 f2) -fill :: Vectorize f => a -> f a -fill x = fmap (const x) empty +-- | partial version of 'devectorize\'' which throws an error +-- if the vector length doesn' match the type length +devectorize :: Vectorize f => V.Vector a -> f a +devectorize x = case devectorize' x of + Right y -> y + Left msg -> error msg +vzipWith :: Vectorize f => (a -> b -> c) -> f a -> f b -> f c +vzipWith f x y = devectorize $ V.zipWith f (vectorize x) (vectorize y) + +vzipWith3 :: Vectorize f => (a -> b -> c -> d) -> f a -> f b -> f c -> f d +vzipWith3 f x y z = devectorize $ V.zipWith3 f (vectorize x) (vectorize y) (vectorize z) + +vzipWith4 :: Vectorize f => (a -> b -> c -> d -> e) -> f a -> f b -> f c -> f d -> f e +vzipWith4 f x y z w = + devectorize $ V.zipWith4 f (vectorize x) (vectorize y) (vectorize z) (vectorize w) + +-- this could me more efficient as a class method, but this is safer +vlength :: Vectorize f => Proxy f -> Int +vlength = V.length . vectorize . (fill () `asFunctorOf`) + where + asFunctorOf :: f a -> Proxy f -> f a + asFunctorOf x _ = x + -- | fmap f == devectorize . (V.map f) . vectorize class Functor f => Vectorize (f :: * -> *) where vectorize :: f a -> V.Vector a - devectorize :: V.Vector a -> f a - empty :: f () + devectorize' :: V.Vector a -> Either String (f a) + fill :: a -> f a default vectorize :: (Generic1 f, GVectorize (Rep1 f)) => f a -> V.Vector a vectorize f = gvectorize (from1 f) - default devectorize :: (Generic1 f, GVectorize (Rep1 f)) => V.Vector a -> f a - devectorize f = to1 (gdevectorize f) - - default empty :: (Generic1 f, GVectorize (Rep1 f)) => f () - empty = to1 gempty - ---vlength :: Vectorize f => Proxy f -> Int ---vlength = const (gvlength (Proxy :: Proxy (Rep1 f))) + default devectorize' :: (Generic1 f, GVectorize (Rep1 f)) => V.Vector a -> Either String (f a) + devectorize' f = fmap to1 (gdevectorize f) + default fill :: (Generic1 f, GVectorize (Rep1 f)) => a -> f a + fill = to1 . gfill -- undecidable, overlapping, orphan instances to get rid of boilerplate + +#if defined(__GLASGOW_HASKELL__) && __GLASGOW_HASKELL__ <= 708 instance Vectorize f => Applicative f where +#else +instance {-# OVERLAPPABLE #-} Vectorize f => Applicative f where +#endif pure = fill x0 <*> x1 = devectorize (V.zipWith id (vectorize x0) (vectorize x1)) + +#if defined(__GLASGOW_HASKELL__) && __GLASGOW_HASKELL__ <= 708 instance Vectorize f => Linear.Additive f where +#else +instance {-# OVERLAPPABLE #-} Vectorize f => Linear.Additive f where +#endif zero = fill 0 + +#if defined(__GLASGOW_HASKELL__) && __GLASGOW_HASKELL__ <= 708 instance Vectorize f => Linear.Metric f where +#else +instance {-# OVERLAPPABLE #-} Vectorize f => Linear.Metric f where +#endif dot x0 x1 = V.sum $ V.zipWith (*) (vectorize x0) (vectorize x1) + +#if defined(__GLASGOW_HASKELL__) && __GLASGOW_HASKELL__ <= 708 instance (Vectorize f, Eq a) => Eq (f a) where +#else +instance {-# OVERLAPPABLE #-} (Vectorize f, Eq a) => Eq (f a) where +#endif x == y = (vectorize x) == (vectorize y) x /= y = (vectorize x) /= (vectorize y) + +#if defined(__GLASGOW_HASKELL__) && __GLASGOW_HASKELL__ <= 708 instance (Vectorize f, Ord a) => Ord (f a) where +#else +instance {-# OVERLAPPABLE #-} (Vectorize f, Ord a) => Ord (f a) where +#endif compare x y = compare (vectorize x) (vectorize y) -instance Vectorize f => Foldable f where + +#if defined(__GLASGOW_HASKELL__) && __GLASGOW_HASKELL__ <= 708 +instance Vectorize f => F.Foldable f where +#else +instance {-# OVERLAPPABLE #-} Vectorize f => F.Foldable f where +#endif foldMap f x = F.foldMap f (vectorize x) foldr f acc0 x = F.foldr f acc0 (vectorize x) -instance Vectorize f => Traversable f where - traverse f x = devectorize <$> T.traverse f (vectorize x) -vlength :: Vectorize f => Proxy f -> Int -vlength = V.length . vectorize . (empty `asFunctorOf`) - where - asFunctorOf :: f a -> Proxy f -> f a - asFunctorOf x _ = x +#if defined(__GLASGOW_HASKELL__) && __GLASGOW_HASKELL__ <= 708 +instance Vectorize f => T.Traversable f where +#else +instance {-# OVERLAPPABLE #-} Vectorize f => T.Traversable f where +#endif + traverse f x = devectorize <$> T.traverse f (vectorize x) class GVectorize (f :: * -> *) where gvectorize :: f a -> V.Vector a - gdevectorize :: V.Vector a -> f a - gempty :: f () + gdevectorize :: V.Vector a -> Either String (f a) + gfill :: a -> f a gvlength :: Proxy f -> Int -vzipWith :: Vectorize f => (a -> b -> c) -> f a -> f b -> f c -vzipWith f x y = devectorize $ V.zipWith f (vectorize x) (vectorize y) - -vzipWith3 :: Vectorize f => (a -> b -> c -> d) -> f a -> f b -> f c -> f d -vzipWith3 f x y z = devectorize $ V.zipWith3 f (vectorize x) (vectorize y) (vectorize z) - -vzipWith4 :: Vectorize f => (a -> b -> c -> d -> e) -> f a -> f b -> f c -> f d -> f e -vzipWith4 f x y z w = - devectorize $ V.zipWith4 f (vectorize x) (vectorize y) (vectorize z) (vectorize w) - -- product type (concatination) instance (GVectorize f, GVectorize g) => GVectorize (f :*: g) where gvectorize (f :*: g) = gvectorize f V.++ gvectorize g gdevectorize v0s | V.length v0s < n0 = - error $ "gdevectorize (f :*: g): V.length v0s < vlength f0 (" ++ - show (V.length v0s) ++ " < " ++ show n0 ++ ")" + Left $ "gdevectorize (f :*: g): V.length v0s < vlength f0 (" ++ + show (V.length v0s) ++ " < " ++ show n0 ++ ")" | V.length v1 /= n1 = - error $ "gdevectorize (f :*: g): V.length v1 /= vlength f1 (" ++ - show (V.length v1) ++ " /= " ++ show n1 ++ ")" - | otherwise = f0 :*: f1 + Left $ "gdevectorize (f :*: g): V.length v1 /= vlength f1 (" ++ + show (V.length v1) ++ " /= " ++ show n1 ++ ")" + | otherwise = case (ef0, ef1) of + (Left msg0, Left msg1) -> + Left $ "gdevectorize (f :*: g): errored on both sides: {" ++ msg0 ++ ", " ++ msg1 ++ "}" + (Left msg0, Right _) -> + Left $ "gdevectorize (f :*: g): errored on left side: " ++ msg0 + (Right _, Left msg1) -> + Left $ "gdevectorize (f :*: g): errored on right side: " ++ msg1 + (Right f0, Right f1) -> Right (f0 :*: f1) where - f0 = gdevectorize v0 - f1 = gdevectorize v1 + ef0 = gdevectorize v0 + ef1 = gdevectorize v1 n0 = gvlength (Proxy :: Proxy f) n1 = gvlength (Proxy :: Proxy g) (v0,v1) = V.splitAt n0 v0s - gempty = gempty :*: gempty + gfill x = gfill x :*: gfill x gvlength = const (nf + ng) where nf = gvlength (Proxy :: Proxy f) @@ -196,8 +247,8 @@ -- Metadata (constructor name, etc) instance GVectorize f => GVectorize (M1 i c f) where gvectorize = gvectorize . unM1 - gdevectorize = M1 . gdevectorize - gempty = M1 gempty + gdevectorize = fmap M1 . gdevectorize + gfill = M1 . gfill gvlength = gvlength . proxy where proxy :: Proxy (M1 i c f) -> Proxy f @@ -207,26 +258,26 @@ instance GVectorize Par1 where gvectorize = V.singleton . unPar1 gdevectorize v = case V.toList v of - [] -> error "gdevectorize Par1: got empty list" - [x] -> Par1 x - xs -> error $ "gdevectorize Par1: got non-1 length: " ++ show (length xs) - gempty = Par1 () + [] -> Left "gdevectorize Par1: got empty list" + [x] -> Right (Par1 x) + xs -> Left $ "gdevectorize Par1: got non-1 length: " ++ show (length xs) + gfill = Par1 gvlength = const 1 -- data with no fields instance GVectorize U1 where gvectorize = const V.empty gdevectorize v - | V.null v = U1 - | otherwise = error $ "gdevectorize U1: got non-null vector, length: " ++ show (V.length v) - gempty = U1 + | V.null v = Right U1 + | otherwise = Left $ "gdevectorize U1: got non-null vector, length: " ++ show (V.length v) + gfill = const U1 gvlength = const 0 -- Constants, additional parameters, and rank-1 recursion instance Vectorize f => GVectorize (Rec1 f) where gvectorize = vectorize . unRec1 - gdevectorize = Rec1 . devectorize - gempty = Rec1 empty + gdevectorize = fmap Rec1 . devectorize' + gfill = Rec1 . fill gvlength = vlength . proxy where proxy :: Proxy (Rec1 f) -> Proxy f @@ -234,17 +285,25 @@ -- composition instance (Vectorize f, GVectorize g) => GVectorize (f :.: g) where - gempty = Comp1 (devectorize (V.replicate k gempty)) + gfill = Comp1 . devectorize'' . V.replicate k . gfill where + devectorize'' x = case devectorize' x of + Right y -> y + Left msg -> error $ "gfill (f :.: g) devectorize error: " ++ msg k = vlength (Proxy :: Proxy f) + gvectorize = V.concatMap gvectorize . vectorize . unComp1 - gdevectorize v = Comp1 (devectorize vs) + gdevectorize v = case partitionEithers (V.toList evs) of + ([], vs) -> fmap Comp1 (devectorize' (V.fromList vs)) + (bad, good) -> Left $ printf "gdevectorize (f :.: g): got %d failures and %d successes" + (length bad) (length good) where kf = vlength (Proxy :: Proxy f) kg = gvlength (Proxy :: Proxy g) - -- vs :: V.Vector (g a) - vs = fmap gdevectorize (splitsAt kg kf v {-:: Vec nf (Vec ng a)-} ) + --evs :: V.Vector (Either String (g a)) + evs = fmap gdevectorize (splitsAt kg kf v {-:: Vec nf (Vec ng a)-} ) + gvlength = const (nf * ng) where nf = vlength (Proxy :: Proxy f) diff --git a/src/Dyno/View/Cov.hs b/src/Dyno/View/Cov.hs --- a/src/Dyno/View/Cov.hs +++ b/src/Dyno/View/Cov.hs @@ -1,6 +1,6 @@ {-# OPTIONS_GHC -Wall -fno-cse #-} -{-# Language ScopedTypeVariables #-} -{-# Language KindSignatures #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE KindSignatures #-} module Dyno.View.Cov ( Cov(..) @@ -12,6 +12,7 @@ , fromMatrix , diag , diag' + , diag'' , nOfVecLen ) where @@ -20,7 +21,7 @@ import qualified Data.Vector as V import qualified Data.Sequence as Seq import System.IO.Unsafe ( unsafePerformIO ) -import qualified Data.Packed.Matrix as Mat +import qualified Numeric.LinearAlgebra as Mat import qualified Casadi.Sparsity as Sparsity import Casadi.Slice ( slice' ) @@ -31,7 +32,7 @@ import Dyno.View.Unsafe.View ( unJ, mkJ ) import Dyno.View.Unsafe.M ( M(UnsafeM), mkM ) -import Dyno.Vectorize ( Vectorize(..) ) +import Dyno.Vectorize ( Vectorize(..), vlength, devectorize ) import Dyno.View.View ( View(..), J ) import Dyno.View.JV ( JV ) import Dyno.View.Viewable ( Viewable(..) ) @@ -83,6 +84,16 @@ where f y k = replicate k offDiag ++ [y] vx = V.toList $ vectorize x + +diag'' :: forall f a . (Vectorize f, Num a) => f a -> f (f a) +diag'' v0 = devectorize $ V.generate n (\k -> devectorize (V.generate n (\j -> gen j k))) + where + v = vectorize v0 + n = vlength (Proxy :: Proxy f) + gen j k + | j /= k = 0 + | otherwise = v V.! k + --data X a = X (J S a) (J S a) deriving (Generic, Show) --instance View X diff --git a/src/Dyno/View/Fun.hs b/src/Dyno/View/Fun.hs --- a/src/Dyno/View/Fun.hs +++ b/src/Dyno/View/Fun.hs @@ -44,7 +44,7 @@ newtype MXFun (f :: * -> *) (g :: * -> *) = MXFun C.MXFunction newtype SXFun (f :: * -> *) (g :: * -> *) = SXFun C.SXFunction -newtype Fun (f :: * -> *) (g :: * -> *) = Fun C.Function +newtype Fun (f :: * -> *) (g :: * -> *) = Fun { unFun :: C.Function } instance Show (MXFun f g) where showsPrec k (MXFun f) = showsPrec k f diff --git a/src/Dyno/View/JV.hs b/src/Dyno/View/JV.hs --- a/src/Dyno/View/JV.hs +++ b/src/Dyno/View/JV.hs @@ -23,7 +23,7 @@ import Dyno.View.View ( View(..), J ) import Dyno.View.Viewable ( Viewable(..) ) -import Dyno.Vectorize ( Vectorize(..), Id, vlength ) +import Dyno.Vectorize ( Vectorize(..), Id, vlength, devectorize ) -- | views into Vectorizable things newtype JV f a = JV { unJV :: f a } deriving (Functor, Generic, Generic1) diff --git a/src/Dyno/View/JVec.hs b/src/Dyno/View/JVec.hs --- a/src/Dyno/View/JVec.hs +++ b/src/Dyno/View/JVec.hs @@ -19,15 +19,16 @@ import Dyno.View.Unsafe.View ( mkJ, unJ ) -import Dyno.TypeVecs ( Vec, unVec, mkVec, reifyVector ) +import Dyno.TypeVecs ( Vec, unVec, reifyVector ) import Dyno.View.Viewable ( Viewable(..) ) import Dyno.View.View ( View(..), J ) +import Dyno.Vectorize ( devectorize ) -- | vectors in View -newtype JVec (n :: k) f a = JVec { unJVec :: Vec n (J f a) } deriving ( Show, Eq ) +newtype JVec (n :: k) f a = JVec { unJVec :: Vec n (J f a) } deriving ( Show ) instance (Dim n, View f) => View (JVec n f) where cat = mkJ . vveccat . fmap unJ . unVec . unJVec - split = JVec . fmap mkJ . mkVec . flip vvertsplit ks . unJ + split = JVec . fmap mkJ . devectorize . flip vvertsplit ks . unJ where ks = V.fromList (take (n+1) [0,m..]) n = reflectDim (Proxy :: Proxy n) @@ -44,7 +45,7 @@ jreplicate' :: forall a n f . (Dim n, View f) => J f a -> JVec n f a jreplicate' el = ret where - ret = JVec (mkVec (V.replicate nvec el)) + ret = JVec (devectorize (V.replicate nvec el)) nvec = reflectDim (Proxy :: Proxy n) jreplicate :: forall a n f . (Dim n, View f, Viewable a) => J f a -> J (JVec n f) a diff --git a/src/Dyno/View/M.hs b/src/Dyno/View/M.hs --- a/src/Dyno/View/M.hs +++ b/src/Dyno/View/M.hs @@ -9,12 +9,17 @@ ( M , sparse, dense , mm + , mv + , vm , ms + , sm , vs + , sv , trans , zeros , eye , diag + , takeDiag , ones , countUp , vsplit @@ -27,12 +32,16 @@ , hcat' , hsplitTup , hsplitTrip + , hsplitQuad , hcatTup , hcatTrip + , hcatQuad , vsplitTup , vsplitTrip + , vsplitQuad , vcatTup , vcatTrip + , vcatQuad , row , col , unrow @@ -51,15 +60,14 @@ import Casadi.CMatrix ( CMatrix ) import Casadi.DMatrix ( DMatrix, dnonzeros, dsparsify ) import qualified Casadi.CMatrix as CM -import qualified Data.Packed.Matrix as HMat -import qualified Numeric.LinearAlgebra.HMatrix as HMat +import qualified Numeric.LinearAlgebra as HMat import Dyno.View.Unsafe.View ( unJ, mkJ ) import Dyno.View.Unsafe.M ( M(UnsafeM), mkM, mkM', unM ) -import Dyno.Vectorize ( Vectorize(..), Id, fill ) +import Dyno.Vectorize ( Vectorize(..), Id, fill, devectorize ) import Dyno.TypeVecs ( Vec, Dim(..) ) -import Dyno.View.View ( View(..), J, JTuple, JTriple ) +import Dyno.View.View ( View(..), J, JTuple, JTriple, JQuad ) import Dyno.View.JV ( JV ) import Dyno.View.JVec ( JVec ) import Dyno.View.Viewable ( Viewable ) @@ -75,12 +83,24 @@ mm :: (View f, View h, CMatrix a) => M f g a -> M g h a -> M f h a mm (UnsafeM m0) (UnsafeM m1) = mkM (CM.mm m0 m1) +mv :: (View f, View g, CMatrix a, Viewable a) => M f g a -> J g a -> J f a +mv m v = uncol $ mm m (col v) + +vm :: (View f, View g, CMatrix a, Viewable a) => J f a -> M f g a -> J g a +vm v m = unrow $ mm (row v) m + ms :: (View f, View h, Viewable a, CMatrix a) => M f g a -> J (JV Id) a -> M f h a -ms (UnsafeM m0) m1 = mkM (m0 * (unJ m1)) +ms m0 m1 = mkM $ (unM m0) * (unJ m1) +sm :: (View f, View h, Viewable a, CMatrix a) => J (JV Id) a -> M f g a -> M f h a +sm m0 m1 = mkM $ (unJ m0) * (unM m1) + vs :: (View f, Viewable a, CMatrix a) => J f a -> J (JV Id) a -> J f a vs m0 m1 = uncol $ ms (col m0) m1 +sv :: (View f, Viewable a, CMatrix a) => J (JV Id) a -> J f a -> J f a +sv m0 m1 = uncol $ sm m0 (col m1) + trans :: (View f, View g, CMatrix a) => M f g a -> M g f a trans (UnsafeM m) = mkM (CM.trans m) @@ -147,6 +167,28 @@ => M f g1 a -> M f g2 a -> M f g3 a -> M f (JTriple g1 g2 g3) a hcatTrip (UnsafeM x) (UnsafeM y) (UnsafeM z) = mkM (CM.horzcat (V.fromList [x,y,z])) +hsplitQuad :: + forall f g0 g1 g2 g3 a . + (View f, View g0, View g1, View g2, View g3, CMatrix a) + => M f (JQuad g0 g1 g2 g3) a -> (M f g0 a, M f g1 a, M f g2 a, M f g3 a) +hsplitQuad (UnsafeM x) = + case V.toList (CM.horzsplit x ncs) of + [g0,g1,g2,g3] -> (mkM g0, mkM g1, mkM g2, mkM g3) + n -> error $ "hsplitQuad made a bad split with length " ++ show (length n) + where + ng0 = size (Proxy :: Proxy g0) + ng1 = size (Proxy :: Proxy g1) + ng2 = size (Proxy :: Proxy g2) + ng3 = size (Proxy :: Proxy g3) + ncs = V.fromList [0,ng0,ng0+ng1,ng0+ng1+ng2,ng0+ng1+ng2+ng3] + +hcatQuad :: + forall f g0 g1 g2 g3 a . + (View f, View g0, View g1, View g2, View g3, CMatrix a) + => M f g0 a -> M f g1 a -> M f g2 a -> M f g3 a -> M f (JQuad g0 g1 g2 g3) a +hcatQuad (UnsafeM x0) (UnsafeM x1) (UnsafeM x2) (UnsafeM x3) = + mkM (CM.horzcat (V.fromList [x0,x1,x2,x3])) + hcat :: forall f g a . (View f, Vectorize g, CMatrix a) @@ -211,6 +253,28 @@ => M f1 h a -> M f2 h a -> M f3 h a -> M (JTriple f1 f2 f3) h a vcatTrip (UnsafeM x) (UnsafeM y) (UnsafeM z) = mkM (CM.vertcat (V.fromList [x,y,z])) +vsplitQuad :: + forall f0 f1 f2 f3 h a . + (View f0, View f1, View f2, View f3, View h, CMatrix a) + => M (JQuad f0 f1 f2 f3) h a -> (M f0 h a, M f1 h a, M f2 h a, M f3 h a) +vsplitQuad (UnsafeM x) = + case V.toList (CM.vertsplit x ncs) of + [f0,f1,f2,f3] -> (mkM f0, mkM f1, mkM f2, mkM f3) + n -> error $ "vsplitQuad made a bad split with length " ++ show (length n) + where + nf0 = size (Proxy :: Proxy f0) + nf1 = size (Proxy :: Proxy f1) + nf2 = size (Proxy :: Proxy f2) + nf3 = size (Proxy :: Proxy f3) + ncs = V.fromList [0,nf0,nf0+nf1,nf0+nf1+nf2,nf0+nf1+nf2+nf3] + +vcatQuad :: + forall f0 f1 f2 f3 h a . + (View f0, View f1, View f2, View f3, View h, CMatrix a) + => M f0 h a -> M f1 h a -> M f2 h a -> M f3 h a -> M (JQuad f0 f1 f2 f3) h a +vcatQuad (UnsafeM x0) (UnsafeM x1) (UnsafeM x2) (UnsafeM x3) = + mkM (CM.vertcat (V.fromList [x0,x1,x2,x3])) + hcat' :: forall f g n a . (View f, View g, Dim n, CMatrix a) @@ -248,6 +312,9 @@ where z = CM.diag (unJ x) +takeDiag :: forall f a . (View f, Viewable a, CMatrix a) => M f f a -> J f a +takeDiag m = mkJ $ CM.diag (unM m) + ones :: forall f g a . (View f, View g, CMatrix a) => M f g a ones = mkM z where @@ -285,19 +352,19 @@ toHMat :: forall n m . (View n, View m) => M n m DMatrix -> HMat.Matrix Double -toHMat (UnsafeM d) = HMat.trans $ (m HMat.>< n) (V.toList v) +toHMat (UnsafeM d) = HMat.tr' $ (m HMat.>< n) (V.toList v) where v = dnonzeros (CM.densify d) n = size (Proxy :: Proxy n) m = size (Proxy :: Proxy m) -fromHMat :: (View f, View g) => HMat.Matrix Double -> M f g DMatrix +fromHMat :: (View f, View g, CMatrix a) => HMat.Matrix Double -> M f g a fromHMat x = case fromHMat' x of Right x' -> x' Left msg -> error msg -fromHMat' :: (View f, View g) => HMat.Matrix Double -> Either String (M f g DMatrix) -fromHMat' = mkM' . CM.vertcat . V.fromList . fmap (CM.trans . CM.fromDVector . V.fromList) . HMat.toLists +fromHMat' :: (View f, View g, CMatrix a) => HMat.Matrix Double -> Either String (M f g a) +fromHMat' = mkM' . CM.fromDMatrix . CM.vertcat . V.fromList . fmap (CM.trans . CM.fromDVector . V.fromList) . HMat.toLists rcond :: (View f, View g) => M f g DMatrix -> Double rcond = HMat.rcond . toHMat diff --git a/src/Dyno/View/Scheme.hs b/src/Dyno/View/Scheme.hs --- a/src/Dyno/View/Scheme.hs +++ b/src/Dyno/View/Scheme.hs @@ -28,9 +28,12 @@ import Dyno.View.Unsafe.M ( unM, mkM' ) import qualified Dyno.View.M as M -import Dyno.View.View ( View(..), J ) +import Dyno.View.View ( View(..), J, JQuad, JTriple, JTuple ) import Dyno.View.Viewable ( Viewable ) +instance (View f0, View f1, View f2, View f3) => Scheme (JQuad f0 f1 f2 f3) +instance (View f0, View f1, View f2) => Scheme (JTriple f0 f1 f2) +instance (View f0, View f1) => Scheme (JTuple f0 f1) class FunctionIO (f :: * -> *) where fromMat :: (CMatrix a, Viewable a) => a -> Either String (f a) diff --git a/src/Dyno/View/Unsafe/View.hs b/src/Dyno/View/Unsafe/View.hs --- a/src/Dyno/View/Unsafe/View.hs +++ b/src/Dyno/View/Unsafe/View.hs @@ -32,10 +32,18 @@ instance (View f, B.Binary a, Viewable a) => B.Binary (J f a) where put = B.put . unJ - get = fmap mkJ B.get + get = do + x <- B.get + case mkJ' x of + Right y -> return y + Left msg -> fail msg instance (View f, S.Serialize a, Viewable a) => S.Serialize (J f a) where put = S.put . unJ - get = fmap mkJ S.get + get = do + x <- S.get + case mkJ' x of + Right y -> return y + Left msg -> fail msg instance Show a => Show (J f a) where showsPrec p (UnsafeJ x) = showsPrec p x diff --git a/src/Dyno/View/View.hs b/src/Dyno/View/View.hs --- a/src/Dyno/View/View.hs +++ b/src/Dyno/View/View.hs @@ -8,7 +8,7 @@ module Dyno.View.View ( View(..) , J - , JNone(..), JTuple(..), JTriple(..) + , JNone(..), JTuple(..), JTriple(..), JQuad(..) , jfill , v2d, d2v , fmapJ, unzipJ @@ -17,8 +17,8 @@ import GHC.Generics ( Generic, Generic1 ) -import Data.Foldable ( Foldable ) -import Data.Traversable ( Traversable ) +import qualified Data.Foldable as F +import qualified Data.Traversable as T import Data.Proxy ( Proxy(..) ) import Data.Vector ( Vector ) import qualified Data.Vector as V @@ -33,13 +33,15 @@ import Dyno.View.Unsafe.View -- some helper types -data JNone a = JNone deriving ( Eq, Generic, Generic1, Show, Functor, Foldable, Traversable ) +data JNone a = JNone deriving ( Eq, Generic, Generic1, Show, Functor, F.Foldable, T.Traversable ) data JTuple f g a = JTuple (J f a) (J g a) deriving ( Generic, Show ) data JTriple f g h a = JTriple (J f a) (J g a) (J h a) deriving ( Generic, Show ) +data JQuad f0 f1 f2 f3 a = JQuad (J f0 a) (J f1 a) (J f2 a) (J f3 a) deriving ( Generic, Show ) instance Vectorize JNone where instance View JNone where instance (View f, View g) => View (JTuple f g) instance (View f, View g, View h) => View (JTriple f g h) +instance (View f0, View f1, View f2, View f3) => View (JQuad f0 f1 f2 f3) jfill :: forall a f . View f => a -> J f (Vector a) jfill x = mkJ (V.replicate n x) diff --git a/tests/Doctests.hs b/tests/Doctests.hs new file mode 100644 --- /dev/null +++ b/tests/Doctests.hs @@ -0,0 +1,10 @@ +{-# OPTIONS_GHC -Wall #-} + +module Main ( main ) where + +import Test.DocTest + +main :: IO () +main = doctest [ "-isrc" + , "src/Dyno/FormatTime.hs" + ] diff --git a/tests/IntegrationTests.hs b/tests/IntegrationTests.hs --- a/tests/IntegrationTests.hs +++ b/tests/IntegrationTests.hs @@ -1,12 +1,12 @@ {-# OPTIONS_GHC -Wall #-} -{-# Language ScopedTypeVariables #-} -{-# Language DataKinds #-} -{-# Language TypeFamilies #-} -{-# Language PolyKinds #-} -{-# Language DeriveFunctor #-} -{-# Language DeriveGeneric #-} -{-# Language RankNTypes #-} -{-# Language FlexibleContexts #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE PolyKinds #-} +{-# LANGUAGE DeriveFunctor #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE FlexibleContexts #-} module IntegrationTests ( integrationTests @@ -23,8 +23,9 @@ import qualified Test.HUnit.Base as HUnit import Test.Framework ( Test, testGroup ) import Test.Framework.Providers.HUnit ( testCase ) +import Linear ( Additive ) -import Dyno.Vectorize ( Vectorize(..), None(..), fill ) +import Dyno.Vectorize ( Vectorize(..), None(..), devectorize, fill ) import Dyno.View.View ( View(..), J ) import Dyno.View.JV ( splitJV ) import Dyno.TypeVecs ( Dim ) @@ -64,10 +65,13 @@ type instance C (IntegrationOcp x p) = x type instance H (IntegrationOcp x p) = None type instance Q (IntegrationOcp x p) = None +type instance QO (IntegrationOcp x p) = None +type instance FP (IntegrationOcp x p) = None +type instance PO (IntegrationOcp x p) = None runIntegration :: forall x p deg n - . ( Vectorize x, Vectorize p, Dim deg, Dim n ) + . ( Vectorize x, Vectorize p, Additive x, Dim deg, Dim n ) => Proxy n -> Proxy deg -> QuadratureRoots -> (forall a . Floating a => x a -> p a -> a -> x a) @@ -75,20 +79,16 @@ -> IO (Either String (x Double)) runIntegration _ _ roots ode x0 p tf = do let ocp :: OcpPhase' (IntegrationOcp x p) - ocp = OcpPhase - { ocpMayer = \_ _ _ _ _ -> 0 - , ocpLagrange = \_ _ _ _ _ _ _ -> 0 - , ocpDae = \x' x _ _ pp t -> ((ode x pp t) `minus` x', None) - , ocpQuadratures = \_ _ _ _ _ _ _ -> None - , ocpBc = \x0' _ _ _ _ -> x0' - , ocpPathC = \_ _ _ _ _ _ -> None - , ocpPathCBnds = None - , ocpBcBnds = fmap (\x -> (Just x, Just x)) x0 - , ocpXbnd = fill (Nothing, Nothing) - , ocpUbnd = None - , ocpZbnd = None - , ocpPbnd = fmap (\x -> (Just x, Just x)) p - , ocpTbnd = (Just tf, Just tf) + ocp = + OcpPhase + { ocpMayer = \_ _ _ _ _ _ -> 0 + , ocpLagrange = \_ _ _ _ _ _ _ _ -> 0 + , ocpDae = \x' x _ _ pp _ t -> ((ode x pp t) `minus` x', None) + , ocpQuadratures = \_ _ _ _ _ _ _ _ -> None + , ocpQuadratureOutputs = \_ _ _ _ _ _ _ _ -> None + , ocpBc = \x0' _ _ _ _ _ -> x0' + , ocpPathC = \_ _ _ _ _ _ _ -> None + , ocpPlotOutputs = \_ _ _ _ _ _ _ _ _ _ _ -> None , ocpObjScale = Nothing , ocpTScale = Nothing , ocpXScale = Nothing @@ -99,9 +99,21 @@ , ocpBcScale = Nothing , ocpPathCScale = Nothing } + ocpInputs :: OcpPhaseInputs' (IntegrationOcp x p) + ocpInputs = + OcpPhaseInputs + { ocpPathCBnds = None + , ocpBcBnds = fmap (\x -> (Just x, Just x)) x0 + , ocpXbnd = fill (Nothing, Nothing) + , ocpUbnd = None + , ocpZbnd = None + , ocpPbnd = fmap (\x -> (Just x, Just x)) p + , ocpTbnd = (Just tf, Just tf) + , ocpFixedP = None + } let guess :: J (CollTraj x None None p n deg) (Vector Double) - guess = cat $ makeGuessSim roots tf x0 (\x _ -> ode x p 0) (\_ _ -> None) p - cp <- makeCollProblem roots ocp guess :: IO (CollProblem x None None p x None x None None n deg) + guess = cat $ makeGuessSim roots tf x0 (\_ x _ -> ode x p 0) (\_ _ -> None) p + cp <- makeCollProblem roots ocp ocpInputs guess :: IO (CollProblem x None None p x None x None None None None None n deg) (msg, opt') <- solveNlp solver (cpNlp cp) Nothing return $ case msg of Left m -> Left m @@ -167,7 +179,7 @@ compareIntegration :: forall x p n deg - . (Vectorize x, Vectorize p, Dim n, Dim deg) + . (Vectorize x, Vectorize p, Additive x, Dim n, Dim deg) => Proxy n -> Proxy deg -> (forall a . Floating a => x a -> p a -> a -> x a) -> x Double -> p Double -> Double -> HUnit.Assertion diff --git a/tests/NewUnitTests.hs b/tests/NewUnitTests.hs --- a/tests/NewUnitTests.hs +++ b/tests/NewUnitTests.hs @@ -2,10 +2,11 @@ module Main ( main ) where -import Data.Monoid ( mempty ) +import qualified Data.Monoid as Mo import Test.Framework ( Test, ColorMode(..), RunnerOptions'(..), TestOptions'(..) , defaultMainWithOpts ) +import QuadratureTests ( quadratureTests ) import VectorizeTests ( vectorizeTests ) import ViewTests ( viewTests ) import IntegrationTests ( integrationTests ) @@ -15,16 +16,22 @@ tests :: [Test] tests = - [ integrationTests + [ quadratureTests + , integrationTests , vectorizeTests , viewTests ] opts :: RunnerOptions' Maybe -opts = mempty { ropt_color_mode = Just ColorAlways - , ropt_threads = Just 1 - , ropt_test_options = Just my_test_opts - } +opts = + Mo.mempty + { ropt_color_mode = Just ColorAlways + , ropt_threads = Just 1 + , ropt_test_options = Just my_test_opts + } my_test_opts :: TestOptions' Maybe -my_test_opts = mempty { topt_timeout = Just (Just 5000000) } +my_test_opts = + Mo.mempty + { topt_timeout = Just (Just 15000000) + } diff --git a/tests/QuadratureTests.hs b/tests/QuadratureTests.hs new file mode 100644 --- /dev/null +++ b/tests/QuadratureTests.hs @@ -0,0 +1,212 @@ +{-# OPTIONS_GHC -Wall #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE DeriveFunctor #-} +{-# LANGUAGE DeriveGeneric #-} + +module QuadratureTests + ( quadratureTests + ) where + +import GHC.Generics ( Generic, Generic1 ) + +import Data.Vector ( Vector ) +import qualified Test.HUnit.Base as HUnit +import Test.Framework ( Test, testGroup ) +import Test.Framework.Providers.HUnit ( testCase ) +import Text.Printf ( printf ) + +import Dyno.Vectorize ( Vectorize(..), None(..), Id(..) ) +import Dyno.View.View ( View(..), J ) +import Dyno.View.JV ( splitJV ) +import Dyno.Solvers +import Dyno.Nlp ( NlpOut(..), Bounds ) +import Dyno.NlpUtils +import Dyno.Ocp +import Dyno.DirectCollocation.Formulate +import Dyno.DirectCollocation.Types +--import Dyno.DirectCollocation.Types ( CollTraj(..) ) +import Dyno.DirectCollocation.Quadratures ( QuadratureRoots(..) ) + + + + +data QuadOcp +type instance X QuadOcp = QuadX +type instance Z QuadOcp = QuadZ +type instance U QuadOcp = QuadU +type instance P QuadOcp = QuadP +type instance R QuadOcp = QuadR +type instance O QuadOcp = QuadO +type instance C QuadOcp = QuadBc +type instance H QuadOcp = None +type instance Q QuadOcp = QuadQ +type instance QO QuadOcp = None +type instance FP QuadOcp = None +type instance PO QuadOcp = None + +data QuadX a = QuadX { xP :: a + , xV :: a + } deriving (Functor, Generic, Generic1, Show) +data QuadZ a = QuadZ deriving (Functor, Generic, Generic1, Show) +data QuadU a = QuadU deriving (Functor, Generic, Generic1, Show) +data QuadP a = QuadP deriving (Functor, Generic, Generic1, Show) +data QuadR a = QuadR (QuadX a) deriving (Functor, Generic, Generic1, Show) +data QuadO a = QuadO a deriving (Functor, Generic, Generic1, Show) +data QuadBc a = QuadBc (QuadX a) deriving (Functor, Generic, Generic1, Show) +data QuadQ a = QuadQ a deriving (Functor, Generic, Generic1, Show) + +instance Vectorize QuadX +instance Vectorize QuadZ +instance Vectorize QuadU +instance Vectorize QuadP +instance Vectorize QuadR +instance Vectorize QuadO +instance Vectorize QuadBc +instance Vectorize QuadQ + +mayer :: Num a => QuadOrLagrange -> a -> QuadX a -> QuadX a -> QuadQ a -> QuadP a -> None a -> a +mayer TestQuadratures _ _ _ (QuadQ qf) _ _ = qf +mayer TestLagrangeTerm _ _ _ _ _ _ = 0 + +data QuadOrLagrange = TestQuadratures | TestLagrangeTerm deriving Show +data StateOrOutput = TestState | TestOutput deriving Show + +lagrange :: Num a => StateOrOutput -> QuadOrLagrange -> QuadX a -> QuadZ a -> QuadU a -> QuadP a -> None a -> QuadO a -> a -> a -> a +lagrange _ TestQuadratures _ _ _ _ _ _ _ _ = 0 +lagrange TestState TestLagrangeTerm (QuadX _ v) _ _ _ _ _ _ _ = v +lagrange TestOutput TestLagrangeTerm _ _ _ _ _ (QuadO v) _ _ = v + +quadratures :: Floating a => + StateOrOutput -> QuadX a -> QuadZ a -> QuadU a -> QuadP a -> None a -> QuadO a -> a -> a -> QuadQ a +quadratures TestState (QuadX _ v) _ _ _ _ _ _ _ = QuadQ v +quadratures TestOutput _ _ _ _ _ (QuadO v) _ _ = QuadQ v + +dae :: Floating a => QuadX a -> QuadX a -> QuadZ a -> QuadU a -> QuadP a -> None a -> a -> (QuadR a, QuadO a) +dae (QuadX p' v') (QuadX _ v) _ _ _ _ _ = (residual, outputs) + where + residual = + QuadR + QuadX { xP = p' - v + , xV = v' - alpha + } + outputs = QuadO v + +alpha :: Fractional a => a +alpha = 7 + +tf :: Fractional a => a +tf = 4.4 + +quadOcp :: StateOrOutput -> QuadOrLagrange -> OcpPhase' QuadOcp +quadOcp stateOrOutput quadOrLag = + OcpPhase + { ocpMayer = mayer quadOrLag + , ocpLagrange = lagrange stateOrOutput quadOrLag + , ocpQuadratures = quadratures stateOrOutput + , ocpQuadratureOutputs = \_ _ _ _ _ _ _ _ -> None + , ocpDae = dae + , ocpBc = bc + , ocpPathC = pathc + , ocpPlotOutputs = \_ _ _ _ _ _ _ _ _ _ _ -> None + , ocpObjScale = Nothing + , ocpTScale = Nothing + , ocpXScale = Nothing + , ocpZScale = Nothing + , ocpUScale = Nothing + , ocpPScale = Nothing + , ocpResidualScale = Nothing + , ocpBcScale = Nothing + , ocpPathCScale = Just None + } + +quadOcpInputs :: OcpPhaseInputs' QuadOcp +quadOcpInputs = + OcpPhaseInputs + { ocpPathCBnds = None + , ocpBcBnds = bcBnds + , ocpXbnd = xbnd + , ocpUbnd = ubnd + , ocpZbnd = QuadZ + , ocpPbnd = QuadP + , ocpTbnd = (Just tf, Just tf) + , ocpFixedP = None + } + +pathc :: Floating a => QuadX a -> QuadZ a -> QuadU a -> QuadP a -> None a -> QuadO a -> a -> None a +pathc _ _ _ _ _ _ _ = None + +xbnd :: QuadX Bounds +xbnd = QuadX { xP = (Nothing, Nothing) + , xV = (Nothing, Nothing) + } + +ubnd :: QuadU Bounds +ubnd = QuadU + +bc :: Floating a => QuadX a -> QuadX a -> QuadQ a -> QuadP a -> None a -> a -> QuadBc a +bc x0 _ _ _ _ _ = QuadBc x0 + +bcBnds :: QuadBc Bounds +bcBnds = + QuadBc + (QuadX + { xP = (Just 0, Just 0) + , xV = (Just 0, Just 0) + }) + +type NCollStages = 120 +type CollDeg = 3 + +guess :: QuadratureRoots -> J (CollTraj' QuadOcp NCollStages CollDeg) (Vector Double) +guess roots = cat $ makeGuess roots tf guessX guessZ guessU parm + where + guessX _ = QuadX { xP = 0 + , xV = 0 + } + guessZ _ = QuadZ + guessU _ = QuadU + parm = QuadP + + + +solver :: Solver +solver = ipoptSolver { options = [ ("expand", Opt True) +-- , ("linear_solver", Opt "ma86") +-- , ("ma86_order", Opt "metis") + , ("print_level", Opt (0 :: Int)) + , ("print_time", Opt False) + ]} + +goodSolution :: NlpOut + (CollTraj QuadX QuadZ QuadU QuadP NCollStages CollDeg) + (CollOcpConstraints QuadX QuadR QuadBc None NCollStages CollDeg) + (Vector Double) + -> HUnit.Assertion +goodSolution out = HUnit.assertBool msg (abs (f - fExpected) < 1e-8 && abs (pF - fExpected) < 1e-8) + where + msg = printf " objective: %.4f, final pos: %.4f, expected: %.4f" f pF fExpected + fExpected = 0.5 * alpha * tf**2 :: Double + QuadX pF _ = splitJV xf' + CollTraj _ _ _ xf' = split (xOpt out) + Id f = splitJV (fOpt out) + +compareIntegration :: (QuadratureRoots, StateOrOutput, QuadOrLagrange) -> HUnit.Assertion +compareIntegration (roots, stateOrOutput, quadOrLag) = HUnit.assert $ do + cp <- makeCollProblem roots (quadOcp stateOrOutput quadOrLag) quadOcpInputs (guess roots) + let nlp = cpNlp cp + (ret, out) <- solveNlp solver nlp Nothing + case ret of + Left msg -> return (HUnit.assertString msg) + Right _ -> return (goodSolution out) :: IO HUnit.Assertion + + +quadratureTests :: Test +quadratureTests = + testGroup "quadrature tests" + [ testCase (show input) (compareIntegration input) + | root <- [Radau, Legendre] + , stateOrOutput <- [TestState, TestOutput] + , quadOrLagr <- [TestQuadratures, TestLagrangeTerm] + , let input = (root, stateOrOutput, quadOrLagr) + ] diff --git a/tests/Utils.hs b/tests/Utils.hs --- a/tests/Utils.hs +++ b/tests/Utils.hs @@ -1,5 +1,5 @@ {-# OPTIONS_GHC -Wall #-} -{-# Language PolyKinds #-} +{-# LANGUAGE PolyKinds #-} module Utils ( reproxy diff --git a/tests/VectorizeTests.hs b/tests/VectorizeTests.hs --- a/tests/VectorizeTests.hs +++ b/tests/VectorizeTests.hs @@ -1,10 +1,11 @@ {-# OPTIONS_GHC -Wall #-} -{-# Language ScopedTypeVariables #-} -{-# Language GADTs #-} -{-# Language DeriveFunctor #-} -{-# Language DeriveGeneric #-} -{-# Language DataKinds #-} -{-# Language PolyKinds #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE DeriveFunctor #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE PolyKinds #-} module VectorizeTests ( Vectorizes(..) @@ -24,6 +25,7 @@ import Test.Framework.Providers.QuickCheck2 ( testProperty ) import Dyno.Vectorize +import Dyno.TypeVecs ( Vec ) import qualified Dyno.TypeVecs as TV import Utils @@ -134,19 +136,43 @@ forall x . (Show (x Int), Eq (x Int), Vectorize x) => Proxy x -> Bool -vectorizeThenDevectorize _ = x0 == x1 +vectorizeThenDevectorize _ = case ex1 of + Right x1 -> x0 == x1 + Left _ -> False where x0 :: x Int x0 = fillInc - x1 :: x Int - x1 = devectorize (vectorize x0) + ex1 :: Either String (x Int) + ex1 = devectorize' (vectorize x0) prop_vecThenDevec :: Vectorizes -> Bool prop_vecThenDevec (Vectorizes _ _ p) = vectorizeThenDevectorize p +transposeUnTranspose :: + forall n m + . (Eq (Vec n (Vec m Int)), Show (Vec n (Vec m Int)), Dim n, Dim m) + => Proxy n -> Proxy m -> Bool +transposeUnTranspose _ _ = x0 == x2 + where + n = TV.reflectDim (Proxy :: Proxy n) + m = TV.reflectDim (Proxy :: Proxy m) + + x0 :: Vec n (Vec m Int) + x0 = TV.mkVec' [TV.mkVec' [(j*m + k) | k <- [0..(m-1)]] | j <- [0..(n-1)]] + + x1 :: Vec m (Vec n Int) + x1 = TV.tvtranspose x0 + + x2 :: Vec n (Vec m Int) + x2 = TV.tvtranspose x1 + +prop_transpose :: Dims -> Dims -> Bool +prop_transpose (Dims _ n) (Dims _ m) = transposeUnTranspose n m + vectorizeTests :: Test vectorizeTests = testGroup "vectorize tests" [ testProperty "vec . devec" prop_vecThenDevec + , testProperty "transposeUnTranspose" prop_transpose ] diff --git a/tests/ViewTests.hs b/tests/ViewTests.hs --- a/tests/ViewTests.hs +++ b/tests/ViewTests.hs @@ -1,9 +1,9 @@ {-# OPTIONS_GHC -Wall -fno-warn-orphans #-} -{-# Language ScopedTypeVariables #-} -{-# Language GADTs #-} -{-# Language DeriveGeneric #-} -{-# Language FlexibleInstances #-} -{-# Language PolyKinds #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE PolyKinds #-} module ViewTests ( Views(..) @@ -17,8 +17,7 @@ import qualified Data.Binary as B import qualified Data.Serialize as S import qualified Data.Traversable as T -import qualified Data.Packed.Matrix as Mat -import qualified Numeric.LinearAlgebra ( ) -- for Eq Matrix +import qualified Numeric.LinearAlgebra as Mat import qualified Data.Vector as V import GHC.Generics ( Generic ) import System.IO.Unsafe ( unsafePerformIO ) @@ -39,7 +38,7 @@ import Dyno.TypeVecs ( Vec, Dim ) import Dyno.Vectorize ( Vectorize(..), Id, fill ) -import Dyno.View.View ( View(..), JNone, JTuple, JTriple ) +import Dyno.View.View ( View(..), JNone, JTuple, JTriple, JQuad ) import Dyno.View.JV ( JV ) import Dyno.View.Viewable ( Viewable ) import Dyno.View.M @@ -454,7 +453,39 @@ m1 = hcatTrip mx my mz return (beEqual m0 m1) +prop_vsplitQuad :: Test +prop_vsplitQuad = + testProperty "vcatQuad . vsplitQuad" $ + \(Views {vwProxy = p0}) (Views {vwProxy = p1}) (Views {vwProxy = p2}) (Views {vwProxy = p3}) (Views {vwProxy = p4}) (CMatrices {cmProxy = p5}) + -> test p0 p1 p2 p3 p4 p5 + where + test :: forall f0 f1 f2 f3 g a + . (View f0, View f1, View f2, View f3, View g, CMatrix a, MyEq a) + => Proxy f0 -> Proxy f1 -> Proxy f2 -> Proxy f3 -> Proxy g -> Proxy a + -> Gen Property + test _ _ _ _ _ _ = do + m0 <- arbitrary :: Gen (M (JQuad f0 f1 f2 f3) g a) + let (mf0,mf1,mf2,mf3) = vsplitQuad m0 + m1 = vcatQuad mf0 mf1 mf2 mf3 + return (beEqual m0 m1) +prop_hsplitQuad :: Test +prop_hsplitQuad = + testProperty "hcatQuad . hsplitQuad" $ + \(Views {vwProxy = p0}) (Views {vwProxy = p1}) (Views {vwProxy = p2}) (Views {vwProxy = p3}) (Views {vwProxy = p4}) (CMatrices {cmProxy = p5}) + -> test p0 p1 p2 p3 p4 p5 + where + test :: forall f g0 g1 g2 g3 a + . (View f, View g0, View g1, View g2, View g3, CMatrix a, MyEq a) + => Proxy f -> Proxy g0 -> Proxy g1 -> Proxy g2 -> Proxy g3 -> Proxy a + -> Gen Property + test _ _ _ _ _ _ = do + m0 <- arbitrary :: Gen (M f (JQuad g0 g1 g2 g3) a) + let (mg0,mg1,mg2,mg3) = hsplitQuad m0 + m1 = hcatQuad mg0 mg1 mg2 mg3 + return (beEqual m0 m1) + + viewTests :: Test viewTests = testGroup "view tests" @@ -473,4 +504,6 @@ , prop_hsplitTup , prop_vsplitTrip , prop_hsplitTrip + , prop_vsplitQuad + , prop_hsplitQuad ]