packages feed

dtmc-0.2.0.0: test/Dtmc/Analysis/FiniteTimeSpec.hs

{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE TypeApplications #-}

module Dtmc.Analysis.FiniteTimeSpec (
    spec,
) where

import Data.Finite (
    Finite,
    finites,
    getFinite,
 )
import Dtmc.Analysis.FiniteTime (
    ConditionalProbabilityError (..),
    Observation (..),
    nStepProbability,
    probability,
    probabilityGiven,
    stepProbability,
 )
import Dtmc.Distribution (
    probabilityAt,
 )
import Dtmc.Distribution.Map qualified as DistributionMap
import Dtmc.Distribution.Vector (
    DistributionVector,
 )
import Dtmc.Distribution.Vector qualified as Vector
import Dtmc.Dynamics (
    evolveVector,
    evolveVectorN,
 )
import Dtmc.State qualified
import Dtmc.TestSupport (
    approxEq,
    chunksOf,
    genSimplexPoint,
    genTransitionRows,
    testTolerance,
 )
import Dtmc.Transition.Kernel qualified as Kernel
import Dtmc.Transition.Matrix (
    TransitionMatrix,
    TransitionMatrixError,
    fromRows,
    power,
    rowAt,
    toRows,
 )
import GHC.Generics (
    Generic,
 )
import Numeric.Natural (
    Natural,
 )
import Test.Hspec (
    Spec,
    describe,
    it,
    shouldBe,
    shouldSatisfy,
 )
import Test.Hspec.QuickCheck (
    prop,
 )
import Test.QuickCheck (
    choose,
    conjoin,
    counterexample,
    forAll,
    property,
    (===),
 )

-- A three-state chain with several impossible one-step transitions.
chain :: TransitionMatrix (Finite 3)
chain =
    either (error . show) id $
        fromRows
            ( chunksOf
                3
                [ 0.5
                , 0.5
                , 0.0
                , 0.0
                , 0.2
                , 0.8
                , 1.0
                , 0.0
                , 0.0
                ]
            )

initial :: DistributionVector (Finite 3)
initial =
    either (error . show) id $
        Vector.fromList [0.6, 0.3, 0.1]

checked :: (Show error) => Either error value -> value
checked = either (error . show) id

asTransitionKernel ::
    (Dtmc.State.FiniteState state) =>
    TransitionMatrix state ->
    Kernel.TransitionKernel state
asTransitionKernel matrix =
    Kernel.fromLaws $ \source ->
        checked $
            DistributionMap.fromList
                [ (destination, stepProbability matrix source destination)
                | destination <- Dtmc.State.finiteStates
                ]

kernelChain :: Kernel.TransitionKernel (Finite 3)
kernelChain = asTransitionKernel chain

mapInitial :: DistributionMap.DistributionMap (Finite 3)
mapInitial =
    checked $
        DistributionMap.fromList
            [ (state, probabilityAt initial state)
            | state <- Dtmc.State.finiteStates
            ]

simpleRandomWalk :: Kernel.TransitionKernel Integer
simpleRandomWalk =
    Kernel.fromLaws $ \state ->
        checked
            (DistributionMap.fromList [(state - 1, 0.5), (state + 1, 0.5)])

closeTo :: Double -> Double -> Bool
closeTo = approxEq testTolerance

{- | Hold for a @Right@ whose 'Double' is within 'testTolerance' of the
expected value; fail for any @Left@ or out-of-tolerance value.
-}
rightCloseTo :: Double -> Either ConditionalProbabilityError Double -> Bool
rightCloseTo expected (Right actual) = approxEq testTolerance actual expected
rightCloseTo _ (Left _) = False

rightResultsClose :: Either error Double -> Either error Double -> Bool
rightResultsClose (Right left) (Right right) = closeTo left right
rightResultsClose (Left _) (Left _) = True
rightResultsClose _ _ = False

data NamedPhase = PhaseA | PhaseB | PhaseC
    deriving (Eq, Ord, Show, Generic)

instance Dtmc.State.FiniteState NamedPhase

namedCycle :: TransitionMatrix NamedPhase
namedCycle =
    either (error . show) id $
        fromRows @NamedPhase
            (chunksOf 3 [0, 1, 0, 0, 0, 1, 1, 0, 0])

twoState :: TransitionMatrix (Finite 2)
twoState =
    either (error . show) id $
        fromRows
            (chunksOf 2 [0.9, 0.1, 0.4, 0.6])

twoStateSquared :: TransitionMatrix (Finite 2)
twoStateSquared =
    either (error . show) id $
        fromRows
            (chunksOf 2 [0.85, 0.15, 0.6, 0.4])

closedFormTransition :: TransitionMatrix (Finite 3)
closedFormTransition =
    either (error . show) id $
        fromRows
            ( chunksOf
                3
                [ 0.1
                , 0.5
                , 0.4
                , 0.1
                , 0.8
                , 0.1
                , 0.0
                , 0.5
                , 0.5
                ]
            )

closedFormProbability :: Int -> Double
closedFormProbability n =
    5 / 63 + 5 / 18 * (0.1 ^ n) - 5 / 14 * (0.3 ^ n)

{- | Five-state transition matrix over states @[A, B, C, D, E]@ used by the
probability examples.
-}
observationMatrix :: TransitionMatrix (Finite 5)
observationMatrix =
    either (error . show) id $
        fromRows
            ( chunksOf
                5
                [ 0
                , 0
                , 0
                , 1
                , 0
                , 1 / 3
                , 0
                , 0
                , 0
                , 2 / 3
                , 0
                , 0
                , 0
                , 0
                , 1
                , 0
                , 0
                , 1 / 3
                , 2 / 3
                , 0
                , 1 / 4
                , 1 / 4
                , 0
                , 0
                , 1 / 2
                ]
            )

-- | Initial law @lambda = [1/4, 1/2, 0, 1/4, 0]@ for the probability examples.
observationInitial :: DistributionVector (Finite 5)
observationInitial =
    either (error . show) id $
        Vector.fromList [1 / 4, 1 / 2, 0, 1 / 4, 0]

spec :: Spec
spec = do
    describe "Observation" $ do
        it "is polymorphic in the state type" $
            (At 2 "rain" :: Observation String) `shouldBe` At 2 "rain"

    describe "stepProbability" $ do
        prop "agrees with rowAt then probabilityAt" $
            forAll (genTransitionRows 3) $ \matrix ->
                case fromRows @(Finite 3) matrix of
                    Right p ->
                        conjoin
                            [ stepProbability p i j
                                === probabilityAt (rowAt p i) j
                            | i <- finites
                            , j <- finites
                            ]
                    Left err ->
                        counterexample
                            ("generated matrix was rejected: " <> show err)
                            False

        it "uses named state constructors" $
            stepProbability namedCycle PhaseB PhaseC
                `shouldBe` 1

    describe "nStepProbability" $ do
        it "is the Kronecker delta at exponent zero" $
            let ijs =
                    [(i, j) | i <- finites, j <- finites] ::
                        [(Finite 2, Finite 2)]
             in map (uncurry (nStepProbability 0 twoState)) ijs
                    `shouldBe` map (\(i, j) -> if i == j then 1 else 0) ijs

        prop "agrees with stepProbability at exponent one" $
            forAll (genTransitionRows 3) $ \matrix ->
                case fromRows @(Finite 3) matrix of
                    Right p ->
                        conjoin
                            [ property $
                                approxEq
                                    testTolerance
                                    (nStepProbability 1 p i j)
                                    (stepProbability p i j)
                            | i <- finites
                            , j <- finites
                            ]
                    Left err ->
                        counterexample
                            ("generated matrix was rejected: " <> show err)
                            False

        it "matches a hand-computed square at exponent two" $
            sequence_
                [ nStepProbability 2 twoState i j
                    `shouldSatisfy` closeTo (probabilityAt (rowAt twoStateSquared i) j)
                | i <- finites :: [Finite 2]
                , j <- finites :: [Finite 2]
                ]

        prop "agrees with the corresponding power entry" $
            forAll (genTransitionRows 3) $ \matrix ->
                case fromRows @(Finite 3) matrix of
                    Right p ->
                        let fourStep = toRows (power 4 p)
                         in conjoin
                                [ property $
                                    approxEq
                                        testTolerance
                                        (nStepProbability 4 p i j)
                                        ( fourStep
                                            !! fromIntegral (getFinite i)
                                            !! fromIntegral (getFinite j)
                                        )
                                | i <- finites
                                , j <- finites
                                ]
                    Left err ->
                        counterexample
                            ("generated matrix was rejected: " <> show err)
                            False

        it "preserves named state types" $
            nStepProbability 2 namedCycle PhaseA PhaseC
                `shouldBe` 1

    describe "nStepProbability hand-computed regressions" $ do
        it "gives P^3(E, D) = 3/8 for the five-state chain" $
            nStepProbability 3 observationMatrix 4 3
                `shouldSatisfy` closeTo (3 / 8)

        it "matches the three-state P^n(2, 0) closed form" $
            mapM_
                ( \n ->
                    nStepProbability n closedFormTransition 2 0
                        `shouldSatisfy` closeTo
                            (closedFormProbability (fromIntegral n))
                )
                ([0, 1, 2, 3, 5, 10, 20] :: [Natural])

    describe "probability for state observations" $ do
        it "returns the initial probability at time zero" $
            conjoin
                [ probability initial chain [At 0 state]
                    === probabilityAt initial state
                | state <- finites
                ]

        prop "agrees with probabilityAt of evolveVectorN"
            $ forAll
                ( (,,)
                    <$> choose (0, 6 :: Int)
                    <*> genSimplexPoint 3
                    <*> genTransitionRows 3
                )
            $ \(k, entries, matrix) ->
                case ( Vector.fromList @(Finite 3) entries
                     , fromRows @(Finite 3) matrix
                     ) of
                    (Right mu, Right p) ->
                        conjoin
                            [ property $
                                approxEq
                                    testTolerance
                                    (probability mu p [At (fromIntegral k) state])
                                    (probabilityAt (evolveVectorN (fromIntegral k) mu p) state)
                            | state <- finites
                            ]
                    result ->
                        counterexample
                            ("generated input was rejected: " <> show result)
                            False

        prop "agrees with repeated evolveVector for small exponents"
            $ forAll
                ( (,,)
                    <$> choose (0, 6 :: Int)
                    <*> genSimplexPoint 3
                    <*> genTransitionRows 3
                )
            $ \(k, entries, matrix) ->
                case ( Vector.fromList @(Finite 3) entries
                     , fromRows @(Finite 3) matrix
                     ) of
                    (Right mu, Right p) ->
                        let iterated = iterate (`evolveVector` p) mu !! k
                         in conjoin
                                [ property $
                                    approxEq
                                        testTolerance
                                        (probability mu p [At (fromIntegral k) state])
                                        (probabilityAt iterated state)
                                | state <- finites
                                ]
                    result ->
                        counterexample
                            ("generated input was rejected: " <> show result)
                            False

    describe "Transition realization independence" $ do
        it "computes transition probabilities on an infinite state type" $ do
            nStepProbability 2 simpleRandomWalk 0 0
                `shouldSatisfy` closeTo 0.5
            nStepProbability 3 simpleRandomWalk 0 0
                `shouldBe` 0

        prop "gives matrices and equivalent kernels the same transition powers" $
            forAll (genTransitionRows 3) $ \rawMatrix ->
                case fromRows rawMatrix ::
                        Either TransitionMatrixError (TransitionMatrix (Finite 3)) of
                    Left problem -> counterexample (show problem) False
                    Right matrix ->
                        let kernel = asTransitionKernel matrix
                         in conjoin
                                [ property $
                                    closeTo
                                        (nStepProbability time matrix source destination)
                                        (nStepProbability time kernel source destination)
                                | source <- finites :: [Finite 3]
                                , destination <- finites :: [Finite 3]
                                , time <- [0 .. 4]
                                ]

        it "matches finite trajectory and observation queries" $ do
            probability mapInitial kernelChain [At 0 0, At 1 1, At 2 2]
                `shouldSatisfy` closeTo
                    (probability initial chain [At 0 0, At 1 1, At 2 2])
            probability
                mapInitial
                kernelChain
                [At 3 2, At 0 0, At 1 1]
                `shouldSatisfy` closeTo
                    (probability initial chain [At 3 2, At 0 0, At 1 1])

        it "matches finite conditional probability queries" $
            rightResultsClose
                (probabilityGiven mapInitial kernelChain [At 2 2] [At 0 0])
                (probabilityGiven initial chain [At 2 2] [At 0 0])
                `shouldBe` True

    describe "probability for consecutive observations" $ do
        it "returns the initial probability for a one-state path" $
            approxEq
                testTolerance
                (probability initial chain [At 0 0])
                (probabilityAt initial 0)
                `shouldBe` True

        it "is lambda_i * P(i, j) for a two-state path" $
            approxEq
                testTolerance
                (probability initial chain [At 0 0, At 1 1])
                (0.6 * 0.5)
                `shouldBe` True

        it "is the product of initial and transition probabilities" $
            approxEq
                testTolerance
                (probability initial chain [At 0 0, At 1 1, At 2 2])
                (0.6 * 0.5 * 0.8)
                `shouldBe` True

        it "is zero for a path with an impossible transition" $ do
            approxEq
                testTolerance
                (probability initial chain [At 0 0, At 1 2])
                0
                `shouldBe` True
            approxEq
                testTolerance
                (probability initial chain [At 0 0, At 1 1, At 2 0])
                0
                `shouldBe` True

        prop "a one-state path equals the initial probability" $
            forAll ((,) <$> genSimplexPoint 3 <*> genTransitionRows 3) $
                \(entries, matrix) ->
                    case ( Vector.fromList @(Finite 3) entries
                         , fromRows @(Finite 3) matrix
                         ) of
                        (Right mu, Right p) ->
                            conjoin
                                [ probability mu p [At 0 i]
                                    === probabilityAt mu i
                                | i <- [0, 1, 2]
                                ]
                        result ->
                            counterexample
                                ("generated input was rejected: " <> show result)
                                False

        prop "a two-state path equals lambda_i * P(i, j)" $
            forAll ((,) <$> genSimplexPoint 3 <*> genTransitionRows 3) $
                \(entries, matrix) ->
                    case ( Vector.fromList @(Finite 3) entries
                         , fromRows @(Finite 3) matrix
                         ) of
                        (Right mu, Right p) ->
                            conjoin
                                [ property $
                                    approxEq
                                        testTolerance
                                        (probability mu p [At 0 i, At 1 j])
                                        ( probabilityAt mu i
                                            * stepProbability p i j
                                        )
                                | i <- [0, 1, 2]
                                , j <- [0, 1, 2]
                                ]
                        result ->
                            counterexample
                                ("generated input was rejected: " <> show result)
                                False

    describe "probability" $ do
        it "returns exactly one for no observations" $
            probability initial chain [] `shouldBe` 1

        it "computes a single state observation" $
            approxEq
                testTolerance
                (probability initial chain [At 1 1])
                (probabilityAt (evolveVectorN 1 initial chain) 1)
                `shouldBe` True

        it "is unchanged by observation order" $
            approxEq
                testTolerance
                (probability initial chain [At 0 0, At 1 1])
                (probability initial chain [At 1 1, At 0 0])
                `shouldBe` True

        it "is unchanged by duplicate observations" $
            approxEq
                testTolerance
                (probability initial chain [At 0 0, At 0 0, At 1 1])
                (probability initial chain [At 0 0, At 1 1])
                `shouldBe` True

        it "is exactly zero for conflicting states at one time" $
            probability initial chain [At 0 0, At 0 1] `shouldBe` 0

        it "agrees with the explicit transition product over times 0, 1, 2" $
            approxEq
                testTolerance
                (probability initial chain [At 0 0, At 1 1, At 2 2])
                (0.6 * 0.5 * 0.8)
                `shouldBe` True

        it "is exactly zero through an impossible transition" $
            probability initial chain [At 0 0, At 1 2] `shouldBe` 0

        it "matches a hand-computed multi-gap example" $
            approxEq
                testTolerance
                ( probability
                    observationInitial
                    observationMatrix
                    [At 2 2, At 3 4, At 6 3]
                )
                (5 / 96)
                `shouldBe` True

        prop "a single observation equals direct evolution" $
            forAll ((,) <$> genSimplexPoint 3 <*> genTransitionRows 3) $
                \(entries, matrix) ->
                    case ( Vector.fromList @(Finite 3) entries
                         , fromRows @(Finite 3) matrix
                         ) of
                        (Right mu, Right p) ->
                            conjoin
                                [ property $
                                    approxEq
                                        testTolerance
                                        (probability mu p [At t i])
                                        (probabilityAt (evolveVectorN t mu p) i)
                                | t <- [0, 1, 2]
                                , i <- [0, 1, 2]
                                ]
                        result ->
                            counterexample
                                ("generated input was rejected: " <> show result)
                                False

        prop "is invariant under observation order" $
            forAll ((,) <$> genSimplexPoint 3 <*> genTransitionRows 3) $
                \(entries, matrix) ->
                    case ( Vector.fromList @(Finite 3) entries
                         , fromRows @(Finite 3) matrix
                         ) of
                        (Right mu, Right p) ->
                            property $
                                approxEq
                                    testTolerance
                                    (probability mu p [At 1 1, At 3 2])
                                    (probability mu p [At 3 2, At 1 1])
                        result ->
                            counterexample
                                ("generated input was rejected: " <> show result)
                                False

    describe "probabilityGiven" $ do
        it "returns the event probability for an empty condition" $
            probabilityGiven initial chain [At 1 1] []
                `shouldSatisfy` rightCloseTo
                    (probability initial chain [At 1 1])

        it "returns one for an empty event and a positive condition" $
            probabilityGiven initial chain [] [At 0 0]
                `shouldSatisfy` rightCloseTo 1

        it "returns one when conditioning an observation on itself" $
            probabilityGiven initial chain [At 1 1] [At 1 1]
                `shouldSatisfy` rightCloseTo 1

        it "ignores observations shared by event and condition" $
            probabilityGiven initial chain [At 1 1] [At 1 1, At 2 2]
                `shouldSatisfy` rightCloseTo 1

        it "returns zero for a conflict against a possible condition" $
            probabilityGiven initial chain [At 1 0] [At 1 1]
                `shouldSatisfy` rightCloseTo 0

        it "reports a zero-probability condition" $
            probabilityGiven initial chain [At 0 0] [At 0 0, At 1 2]
                `shouldBe` Left ZeroProbabilityCondition

        it "reports a contradictory condition" $
            probabilityGiven initial chain [At 0 0] [At 1 1, At 1 2]
                `shouldBe` Left ZeroProbabilityCondition

        it "is unaffected by event and condition ordering" $ do
            probabilityGiven initial chain [At 2 2, At 1 1] [At 0 0]
                `shouldSatisfy` rightCloseTo 0.4
            probabilityGiven initial chain [At 1 1, At 2 2] [At 0 0]
                `shouldSatisfy` rightCloseTo 0.4

    describe "probabilityGiven hand-computed regressions" $ do
        it "gives P(X10=D, X11=D | X3=A, X7=E) = 1/4" $
            probabilityGiven
                observationInitial
                observationMatrix
                [At 10 3, At 11 3]
                [At 3 0, At 7 4]
                `shouldSatisfy` rightCloseTo (1 / 4)

        it "accepts an out-of-order event and gives 15/92" $
            probabilityGiven
                observationInitial
                observationMatrix
                [At 6 3, At 2 2]
                [At 3 4]
                `shouldSatisfy` rightCloseTo (15 / 92)

        it "gives P(X2=C) = 5/36" $
            approxEq
                testTolerance
                (probability observationInitial observationMatrix [At 2 2])
                (5 / 36)
                `shouldBe` True

        it "gives P(X3=E) = 23/72" $
            approxEq
                testTolerance
                (probability observationInitial observationMatrix [At 3 4])
                (23 / 72)
                `shouldBe` True

        it "gives P(X2=C, X3=E, X6=D) = 5/96" $
            approxEq
                testTolerance
                ( probability
                    observationInitial
                    observationMatrix
                    [At 2 2, At 3 4, At 6 3]
                )
                (5 / 96)
                `shouldBe` True