{-# OPTIONS_GHC -Wall #-} {-# Language ScopedTypeVariables #-} {-# Language DeriveGeneric #-} {-# Language PolyKinds #-} module Dyno.DirectCollocation.Dynamic ( DynPlotPoints , CollTrajMeta(..) , addCollocationChannel , toMeta , toMetaCov , dynPlotPoints , catDynPlotPoints , NameTree(..) ) where import GHC.Generics ( Generic ) import Data.Proxy ( Proxy(..) ) import Data.List ( mapAccumL ) import Data.Tree ( Tree(..) ) 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.Tree as Tree import Data.Binary ( Binary ) import Linear.V import Accessors ( AccessorTree(..), Lookup(..), accessors ) import PlotHo ( Plotter, addChannel ) import Dyno.View.Unsafe.View ( unJ, unJ' ) import Dyno.Vectorize ( Vectorize, Id(..), fill ) import Dyno.View.JV ( JV, splitJV ) import Dyno.View.View ( View(..), J ) import Dyno.View.JVec ( JVec(..) ) import qualified Dyno.TypeVecs as TV import Dyno.TypeVecs ( Vec ) 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 where toSignalTree :: (DynPlotPoints Double, CollTrajMeta) -> [Tree ( String , String , Maybe ((DynPlotPoints Double, CollTrajMeta) -> [[(Double, Double)]]) )] toSignalTree = forestFromMeta . snd 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 ] data DynPlotPoints a = DynPlotPoints (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 instance Binary a => Binary (DynPlotPoints a) 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) dynPlotPoints :: forall x z u p o n deg a . (Dim n, Dim deg, Real a, Fractional a, Show a, Vectorize x, Vectorize z, Vectorize u, Vectorize o, Vectorize p) => 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 x) (Vector a)) -> DynPlotPoints a dynPlotPoints quadratureRoots (CollTraj tf' _ stages' xf) outputs = DynPlotPoints xss' zss uss oss xdss where nStages = size (Proxy :: Proxy (JVec n (JV Id))) tf,h :: a Id tf = splitJV tf' h = tf / fromIntegral nStages taus :: Vec deg a taus = mkTaus quadratureRoots 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,zss,uss,oss,xdss :: Vector (Vector (a, Vector a)) (xss,zss,uss,oss,xdss) = V.unzip5 xzuoxds -- todo: check this final time tf'' against expected tf (_tf'', xzuoxds) = T.mapAccumL f 0 $ V.zip (TV.unVec stages) (TV.unVec outputs) 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 x) (Vector a)) ) -> ( 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', (xdos, xnext)) = (tnext, (xs,zs,us,os,xds)) 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',zs,us,os,xds :: Vector (a, Vector a) (xs',zs,us,os,xds) = V.unzip5 $ TV.unVec $ TV.tvzipWith3 g xzus0 xdos taus g (CollPoint x z u) (o,x') tau = ( (t,unJ' "x" x) , (t,unJ' "z" z) , (t,unJ' "u" u) , (t,unJ' "o" o) , (t,unJ' "x'" x') ) where t = t0 + h*tau data NameTree = NameTreeNode (String,String) [(String,NameTree)] | NameTreeLeaf Int deriving (Show, Eq, Generic) instance Binary NameTree data CollTrajMeta = CollTrajMeta { ctmX :: NameTree , ctmZ :: NameTree , ctmU :: NameTree , ctmP :: NameTree , ctmO :: NameTree , ctmQ :: NameTree } deriving (Eq, Generic, Show) instance Binary CollTrajMeta namesFromAccTree :: AccessorTree a -> NameTree namesFromAccTree x = (\(_,(_,y)) -> y) $ namesFromAccTree' 0 ("",x) namesFromAccTree' :: Int -> (String, AccessorTree a) -> (Int, (String, NameTree)) namesFromAccTree' k (nm, ATGetter _) = (k+1, (nm, NameTreeLeaf k)) namesFromAccTree' k0 (nm, Data names ats) = (k, (nm, NameTreeNode names children)) where (k, children) = mapAccumL namesFromAccTree' k0 ats type MetaTree a = Tree.Forest (String, String, Maybe ((DynPlotPoints a, CollTrajMeta) -> [[(a,a)]])) forestFromMeta :: CollTrajMeta -> MetaTree Double forestFromMeta meta = [xTree,zTree,uTree,oTree,xdTree] 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) blah :: forall f c t . (Functor f, F.Foldable f) => (c -> f (f (t, Vector t))) -> String -> NameTree -> Tree (String, String, Maybe ((c,CollTrajMeta) -> [[(t, t)]])) blah f myname (NameTreeNode (nm1,_) children) = Tree.Node (myname,nm1,Nothing) $ map (uncurry (blah f)) children blah f myname (NameTreeLeaf k) = Tree.Node (myname,"",Just (woo . f . fst)) [] where woo :: f (f (t, Vector t)) -> [[(t, t)]] woo = F.toList . fmap (F.toList . fmap (\(t,x) -> (t, x V.! k))) toMeta :: forall x z u p o q n deg . (Lookup (x ()), Lookup (z ()), Lookup (u ()), Lookup (p ()), Lookup (o ()), Lookup (q ()), Vectorize x, Vectorize z, Vectorize u, Vectorize p, Vectorize o, Vectorize q, Dim n, Dim deg) => Proxy o -> Proxy q -> Proxy (CollTraj x z u p n deg) -> CollTrajMeta toMeta _ _ _ = CollTrajMeta { ctmX = namesFromAccTree $ accessors (fill () :: x ()) , ctmZ = namesFromAccTree $ accessors (fill () :: z ()) , ctmU = namesFromAccTree $ accessors (fill () :: u ()) , ctmP = namesFromAccTree $ accessors (fill () :: p ()) , ctmO = namesFromAccTree $ accessors (fill () :: o ()) , ctmQ = namesFromAccTree $ accessors (fill () :: q ()) } toMetaCov :: forall sx x z u p o q n deg . (Lookup (x ()), Lookup (z ()), Lookup (u ()), Lookup (p ()), Lookup (o ()), Lookup (q ()), Vectorize x, Vectorize z, Vectorize u, Vectorize p, Vectorize o, Vectorize q, Dim n, Dim deg) => Proxy o -> Proxy q -> Proxy (CollTrajCov sx x z u p n deg) -> CollTrajMeta toMetaCov po pq _ = toMeta po pq (Proxy :: Proxy (CollTraj x z u p n deg))