Mdp.hs 13 KB
Newer Older
1 2 3
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE MultiParamTypeClasses #-}
4 5
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE OverloadedStrings #-}
6
{-# LANGUAGE TemplateHaskell #-}
7
{-# LANGUAGE FlexibleInstances #-}
8

9
module Mdp (mdpP, mdpB, valmariMdpB, valmariMdp2B) where
10

11
import           Data.List                      ( intersperse )
12 13
import           Data.Foldable
import           Control.Applicative
14
import           Data.Ord
15

16 17 18 19 20 21 22 23
import           Data.IntMap                    ( IntMap )
import qualified Data.IntMap                   as M
import qualified Data.Text.Lazy.Builder        as Build
import qualified Data.Text.Lazy.Builder.Int    as Build
import qualified Data.Text.Lazy.Builder.RealFloat
                                               as Build
import           Data.Vector                    ( Vector )
import qualified Data.Vector                   as V
24
import qualified Data.Vector.Algorithms.Intro  as V
25
import           Data.Text                      ( Text )
26
import qualified Data.Text                     as T
27
import           Lens.Micro.Platform
28 29

import           Parser
30
import           StatesFile                     ( Partition(..) )
31 32 33 34 35 36

----------------------------------------------------------------------
-- Types
----------------------------------------------------------------------

data Mdp = Mdp
37 38 39
  { mdpNumStates :: Int
  , mdpNumChoices :: Int
  , mdpTransitions :: Vector Transition
40
  } deriving (Show)
41 42 43

-- TODO Support optional action label
data Transition = Transition
44 45 46 47 48
  { transitionSource :: Int
  , transitionChoice :: Int
  , transitionTarget :: Int
  , transitionProbability :: Double
  , transitionAction :: Maybe Text
49
  } deriving (Show)
50

51 52
makeFields ''Mdp
makeFields ''Transition
53

54 55 56 57 58 59 60 61 62 63 64 65 66
----------------------------------------------------------------------
-- Parser
----------------------------------------------------------------------

mdpP :: Parser Mdp
mdpP = Mdp <$> decimalP <*> decimalP <*> transitionsP

transitionsP :: Parser (Vector Transition)
transitionsP = do
  numTrans <- decimalP
  V.replicateM numTrans transitionP

transitionP :: Parser Transition
67 68 69 70 71 72 73
transitionP =
  Transition
    <$> decimalP
    <*> decimalP
    <*> decimalP
    <*> doubleP
    <*> optional nameP
74 75 76 77 78

----------------------------------------------------------------------
-- Builder
----------------------------------------------------------------------

79 80
mdpB :: Maybe Partition -> Mdp -> Build.Builder
mdpB partition mdp =
81
  functorB partition <> "\n" <> transitionsB partition (mdp ^. transitions)
82

83
functorB :: Maybe Partition -> Build.Builder
84
functorB Nothing = "P(Nx(DX))\n"
85
functorB (Just partition) =
86
  (Build.decimal (numBlocks partition)) <> " x P(Nx(DX))\n"
87

88 89
transitionsB :: Maybe Partition -> Vector Transition -> Build.Builder
transitionsB partition ts =
90
  let stateMap :: IntMap [Transition] =
91 92
        V.foldl' (\m t -> M.insertWith (++) (t ^. source) [t] m) M.empty ts
  in  foldMap (uncurry (choicesB partition)) (M.toList stateMap)
93

94
choicesB :: Maybe Partition -> Int -> [Transition] -> Build.Builder
95 96 97 98 99 100 101 102 103 104 105 106
choicesB partition source ts =
  let choiceMap :: IntMap [Transition] =
        foldl' (\m t -> M.insertWith (++) (t ^. choice) [t] m) M.empty ts
  in  stateB source
      <> ": "
      <> partStart partition
      <> "{"
      <> fold
           (intersperse ", " (map (uncurry transitionB) (M.toList choiceMap)))
      <> "}"
      <> partEnd partition
      <> "\n"
107 108 109
  where
    partStart Nothing = ""
    partStart (Just part) =
110 111
      "(" <> Build.decimal (stateAssignment part V.! source) <> ", "
    partEnd Nothing  = ""
112
    partEnd (Just _) = ")"
113

114 115 116 117 118 119 120
transitionB :: Int -> [Transition] -> Build.Builder
transitionB choice successors =
  "("
    <> Build.decimal choice
    <> ", {"
    <> fold (intersperse ", " (map successorB successors))
    <> "})"
121 122
  where
    successorB :: Transition -> Build.Builder
123 124 125
    successorB t = stateB (t ^. target) <> ": " <> Build.formatRealFloat
      Build.Fixed
      Nothing
126
      (t ^. probability)
127 128

stateB :: Int -> Build.Builder
129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149
stateB = ("s" <>) . Build.decimal

----------------------------------------------------------------------
-- Builder for Valmari format
--
-- This format is essentially:
--
-- MC              ::=  <sizes> <l_transitions> <w_transitions> <blocks>
-- <sizes>         ::= num_states num_choices num_l_trans num_w_trans num_blocks
-- <l_transitions> ::= <l_transition>*
-- <l_transition>  ::= source choice target
-- <w_transitions> ::= <w_transition>*
-- <w_transition>  ::= source weight target
-- <blocks>        ::= <block>*
-- <block>         ::= state* 0
--
-- Where blocks are the blocks of the initial partition. Also, one block is
-- omitted.
----------------------------------------------------------------------

data ValmariMDP = ValmariMDP
150 151 152 153 154
  { valmariMDPNumStates :: Int
  , valmariMDPNumActionStates :: Int
  , valmariMDPNumChoices :: Int
  , valmariMDPChoices :: Vector ValmariChoice
  , valmariMDPProbabilities :: Vector ValmariTransition
155
  } deriving (Show)
156

157 158 159 160
data ValmariChoice = ValmariChoice
  { _valmariChoiceSource :: Int
  , _valmariChoiceChoice :: Int
  , _valmariChoiceTarget :: Int
161
  } deriving (Show)
162 163


164 165
-- From an action state to a normal state
data ValmariTransition = ValmariTransition
166 167 168
  { _valmariTransitionSource :: Int
  , _valmariTransitionProbability :: Double
  , _valmariTransitionTarget :: Int
169
  } deriving (Show)
170

171 172 173
makeFields ''ValmariMDP
makeFields ''ValmariChoice
makeFields ''ValmariTransition
174 175 176 177 178

type ValmariPartition = [[Int]]


convertToValmari :: Maybe Partition -> Mdp -> (ValmariMDP, ValmariPartition)
179
convertToValmari maybePart mdp =
180
  let nums    = mdp ^. numStates
181
      (numa, transMap) = mkValmariMap mdp
182 183 184
      numc    = mdp ^. numChoices
      actions = V.fromList (concatMap mkChoices (M.toList transMap))
      probs   = V.fromList (concatMap mkTransitions (transMap ^.. each . each))
185 186 187
      actionStates = map (+ nums) [0 .. numa - 1]
      normalStates = [0 .. nums - 1]
      part    = actionStates : maybe [normalStates] partToBlocks maybePart
188 189 190 191 192 193
  in  ( ValmariMDP { valmariMDPNumStates     = nums
                   , valmariMDPNumActionStates = numa
                   , valmariMDPChoices       = actions
                   , valmariMDPProbabilities = probs
                   , valmariMDPNumChoices    = numc
                   }
194 195
      , part
      )
196 197 198 199
  where
    mkChoices (s, cs) = map (mkChoice s) cs
    mkChoice s (t, c) = ValmariChoice s (fst c) t
    mkTransitions (s, (_, ts)) = map (uncurry (ValmariTransition s)) ts
200

201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232

-- | Alternative method of modeling the same coalgebra. Here, we don't use
-- labels but partition the Distribution states on the label that leads to them.
-- This models the functor P(AxDX) more faithfully.
convertToValmari2 :: Maybe Partition -> Mdp -> (ValmariMDP, ValmariPartition)
convertToValmari2 maybePart mdp =
  let
    nums    = mdp ^. numStates
    (numa, transMap) = mkValmariMap mdp
    numc    = mdp ^. numChoices
    actions = V.fromList (concatMap mkChoices (M.toList transMap))
    probs   = V.fromList (concatMap mkTransitions (transMap ^.. each . each))
    stateByChoice = M.fromListWith
      (++)
      (map (\x -> (x ^. choice, [(x ^. target) + nums])) (V.toList actions))
    part = (maybe [[0 .. nums - 1]] partToBlocks maybePart)
      ++ M.elems stateByChoice
  in
    ( ValmariMDP { valmariMDPNumStates     = nums
                 , valmariMDPNumActionStates = numa
                 , valmariMDPChoices       = actions & each . choice .~ 0
                 , valmariMDPProbabilities = probs
                 , valmariMDPNumChoices    = 1
                 }
    , part
    )
  where
    mkChoices (s, cs) = map (mkChoice s) cs
    mkChoice s (t, c) = ValmariChoice s (fst c) t
    mkTransitions (s, (_, ts)) = map (uncurry (ValmariTransition s)) ts


233
-- State -> Choices
234 235
-- each choice is (choice state, choice number, [(probability, successor)])
type ValmariMap = IntMap [(Int, (Int, [(Double, Int)]))]
236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251


-- | Returns the total number of action states and the valmari map
mkValmariMap :: Mdp -> (Int, ValmariMap)
mkValmariMap mdp = M.mapAccum uniquifyChoices
                              0
                              (V.foldl' ins M.empty (mdp ^. transitions))
  where
    ins
      :: IntMap (IntMap [(Double, Int)])
      -> Transition
      -> IntMap (IntMap [(Double, Int)])
    ins m t =
      m
        &  (at (t ^. source) . non M.empty)
        .  (at (t ^. choice) . non [])
252
        %~ ((t ^. probability, t ^. target) :)
253 254

    uniquifyChoices
255
      :: Int -> IntMap [(Double, Int)] -> (Int, [(Int, (Int, [(Double, Int)]))])
256
    uniquifyChoices current actions =
257
      let newActions = zip [current ..] (M.toList actions)
258 259 260 261 262 263 264 265 266 267 268 269 270 271
      in  (current + length newActions, newActions)


partToBlocks :: Partition -> [[Int]]
partToBlocks p =
  let blockMap = V.ifoldl' (\m s b -> m & at b . non [] %~ (s :))
                           M.empty
                           (stateAssignment p)
  in  M.elems blockMap


valmariMdpB :: Maybe Partition -> Mdp -> Build.Builder
valmariMdpB partition mdp =
  let (vmdp, vpartition) = convertToValmari partition mdp
272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289
  in  valmariMdpBImpl vmdp vpartition


valmariMdp2B :: Maybe Partition -> Mdp -> Build.Builder
valmariMdp2B partition mdp =
  let (vmdp, vpartition) = convertToValmari2 partition mdp
  in  valmariMdpBImpl vmdp vpartition


valmariMdpBImpl :: ValmariMDP -> ValmariPartition -> Build.Builder
valmariMdpBImpl vmdp vpartition =
  valmariSizesB vmdp vpartition
    <> "\n"
    <> valmariLTransitions vmdp
    <> "\n"
    <> valmariWTransitions vmdp
    <> "\n"
    <> valmariBlocks vpartition
290 291 292 293


valmariSizesB :: ValmariMDP -> ValmariPartition -> Build.Builder
valmariSizesB mdp part =
294
  Build.decimal (mdp ^. numStates + mdp ^. numActionStates)
295
    <> " "
296
    <> Build.decimal (mdp ^. numChoices)
297
    <> " "
298
    <> Build.decimal (length (mdp ^. choices))
299
    <> " "
300
    <> Build.decimal (length (mdp ^. probabilities))
301
    <> " "
302
    <> Build.decimal (length part)
303 304 305


valmariLTransitions :: ValmariMDP -> Build.Builder
306
valmariLTransitions mdp = foldMap vChoice (mdp ^.. choices . each)
307
  where
308 309 310
    vChoice :: ValmariChoice -> Build.Builder
    vChoice c =
      vState (c ^. source)
311
        <> " "
312
        <> Build.decimal (c ^. choice + 1)
313
        <> " "
314
        <> vActionState mdp (c ^. target)
315 316 317
        <> "\n"

valmariWTransitions :: ValmariMDP -> Build.Builder
318
valmariWTransitions mdp = foldMap vTrans (mdp ^.. probabilities . each)
319
  where
320 321 322
    vTrans :: ValmariTransition -> Build.Builder
    vTrans t =
      vActionState mdp (t ^. source)
323
        <> " "
324
        <> Build.formatRealFloat Build.Fixed Nothing (t ^. probability)
325
        <> " "
326
        <> vState (t ^. target)
327 328
        <> "\n"

329

330
vState :: Int -> Build.Builder
331 332
vState = Build.decimal . (+ 1)

333 334

vActionState :: ValmariMDP -> Int -> Build.Builder
335
vActionState mdp = Build.decimal . (+ (mdp ^. numStates)) . (+ 1)
336 337 338


valmariBlocks :: [[Int]] -> Build.Builder
339 340
valmariBlocks part =
  part
341
    & tail
342 343 344 345 346
    & (each . each %~ vState)
    & (each %~ intersperse " ")
    & (each %~ fold)
    & (each %~ (<> " 0\n"))
    & fold
347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417

----------------------------------------------------------------------
-- Builder for .aut-Format used in mcrl2
--
-- This format is informally:
--
-- des (initial distribution, nr transitions, nr states)
-- (state, "label", out distribution)
-- ...
----------------------------------------------------------------------

data Mcrl2Mdp = Mcrl2Mdp
  { mcrl2MdpNumStates :: Int
  , mcrl2MdpTransitions :: Vector Mcrl2Transition
  , mcrl2MdpOutDistribution :: Mcrl2Distribution
  } deriving (Show)

data Mcrl2Transition = Mcrl2Transition
  { mcrl2TransitionSource :: Int
  , mcrl2TransitionLabel :: Text
  , mcrl2TransitionOutDistribution :: Mcrl2Distribution
  } deriving (Show)

type Mcrl2Distribution = Vector Mcrl2PropTrans

data Mcrl2PropTrans = Mcrl2PropTrans
  { mcrl2PropTransProbability :: Double
  , mcrl2PropTransTarget :: Int
  } deriving (Show)

makeFields ''Mcrl2Mdp
makeFields ''Mcrl2Transition
makeFields ''Mcrl2PropTrans


-- TODO Fix initial state: For now the initial state is always 0, but that
-- doesn't correspond to the actual prism model. Maybe we should provide a
-- uniform initial distribution over all states.

convertToMcrl :: Mdp -> Mcrl2Mdp
convertToMcrl mdp = Mcrl2Mdp (mdp ^. numStates) (convertTransitionsToMcrl (mdp^.transitions)) (pure initEdge)
  where
    initEdge = Mcrl2PropTrans 1.0 0


convertTransitionsToMcrl :: Vector Transition -> Vector Mcrl2Transition
convertTransitionsToMcrl trans = mkMcrl2Tans <$> partitionVector (view source) sorted
  where
    sorted = V.create $ do
      v <- V.thaw trans -- TODO Maybe use unsafe thaw
      V.sortBy (comparing $ view source) v
      return v

    shouldGroup t = (t^.source, t^.choice)


mkMcrl2Tans :: Vector Transition -> Mcrl2Transition
mkMcrl2Tans distri = Mcrl2Transition (h ^. source) (T.pack $ show $ h ^. choice) (mkPropTrans <$> distri)
  where h = V.head distri


mkPropTrans :: Transition -> Mcrl2PropTrans
mkPropTrans t = Mcrl2PropTrans (t ^. probability) (t ^. target)


partitionVector :: Eq b => (a -> b) -> Vector a -> Vector (Vector a)
partitionVector pred = V.unfoldr helper
  where
    helper vec
      | null vec  = Nothing
      | otherwise = Just $ V.span (((pred (V.head vec)) ==) . pred) vec