Internal.hs

{- |
Module      : Dtmc.Transition.Matrix.Internal
Description : Raw carrier for transition matrices (unsafe underbelly).

Raw carrier behind t'Dtmc.Transition.Matrix.TransitionMatrix': an hmatrix
matrix paired with its lazy support graph. The public smart constructor
validates its square shape and canonicalises rows; this internal module
exposes unchecked construction.

The constructor is positional so the public matrix projection cannot act as a
record-update setter and desynchronise the matrix from its cached graph.
-}
module Dtmc.Transition.Matrix.Internal (
    TransitionMatrix (TransitionMatrix),
    unTransitionMatrix,
    tmSupport,
    unsafeTransitionMatrix,
    matrixRowAt,
) where

import Dtmc.Distribution.Map (
    fromDistribution,
 )
import Dtmc.Distribution.Vector.Internal (
    DistributionVector (DistributionVector),
 )
import Dtmc.State (
    FiniteState,
 )
import Dtmc.State.Internal (
    stateCardinalityInt,
    stateIndexInt,
 )
import Dtmc.Transition (
    Transition (..),
 )
import Dtmc.Transition.Matrix.Internal.Graph (
    Graph,
    fromAdjacency,
 )
import Numeric.LinearAlgebra qualified as LA

{- | A stored square matrix whose rows and columns follow the canonical order
of its finite state type. Entry @(i,j)@ is the transition probability from
state @i@ to state @j@. 'Dtmc.Transition.Matrix.fromKernel' materialises
already-validated rows, while 'Dtmc.Transition.Matrix.fromRows' applies
tolerant row validation and canonicalisation. The internal constructor and
arithmetic instances do not revalidate.

Each value also carries its support graph as a /lazy/ second argument, so any
graph-based analyses on the same value share one build. Construct internal
values with @unsafeTransitionMatrix@ rather than pairing a matrix and graph
directly.
-}
data TransitionMatrix state
    = -- | Unchecked matrix/cache pair; the graph must match the matrix.
      TransitionMatrix (LA.Matrix Double) Graph

-- Nominal role prevents coercion between distinct state types, including
-- state types with the same cardinality.
type role TransitionMatrix nominal

{- | Return the stored matrix unchanged without forcing the support graph.

Complexity: @O(1)@ time and @O(1)@ space.
-}
unTransitionMatrix ::
    TransitionMatrix state ->
    LA.Matrix Double
unTransitionMatrix (TransitionMatrix matrix _) = matrix

{- | Return the lazy support graph, with edge @i -> j@ exactly when the stored
entry is strictly positive. No tolerance is applied: a tiny positive rounding
value creates an edge, while zero or a negative value does not.

The result is shared by later analyses of the same value.

Complexity: @O(1)@ projection time and @O(1)@ projection space. The first
analysis that forces the graph takes @O(n^2)@ time and @O(n^2)@ temporary
space; the resulting graph occupies @O(n + E)@ space for @E@ support edges.
-}
tmSupport :: TransitionMatrix state -> Graph
tmSupport (TransitionMatrix _ support) = support

-- Manual 'Show': 'Graph' has no 'Show', and the derived cache should not
-- appear in the rendering.
instance Show (TransitionMatrix state) where
    showsPrec d p =
        showParen (d > 10) $
            showString "TransitionMatrix "
                . showsPrec 11 (unTransitionMatrix p)

{- | Pair a raw matrix with its lazy support graph. This performs no
row-stochastic, finiteness, or simplex validation; internal callers must
establish the required invariant.

Complexity: @O(1)@ construction time and @O(1)@ construction space. Forcing
the support graph takes @O(n^2)@ time and @O(n^2)@ temporary space; the graph
occupies @O(n + E)@ space for @E@ support edges.
-}
unsafeTransitionMatrix ::
    LA.Matrix Double ->
    TransitionMatrix state
unsafeTransitionMatrix matrix =
    TransitionMatrix matrix (supportGraphOf matrix)

{- | Wrap one stored matrix row as a distribution vector without revalidation.
The finite-state index makes the lookup total.

Complexity: excluding 'Dtmc.State.stateIndex', @O(n)@ time and @O(n)@ result
space for state cardinality @n@.
-}
matrixRowAt ::
    (FiniteState state) =>
    TransitionMatrix state ->
    state ->
    DistributionVector state
matrixRowAt matrix state = DistributionVector row
  where
    stored = unTransitionMatrix matrix
    row =
        LA.flatten
            ( LA.subMatrix
                (stateIndexInt state, 0)
                (1, LA.cols stored)
                stored
            )

instance (FiniteState state) => Transition (TransitionMatrix state) where
    type TransitionState (TransitionMatrix state) = state

    transitionLaw matrix =
        fromDistribution . matrixRowAt matrix

-- Use strict positivity without tolerance so graph queries reflect the stored
-- matrix exactly; keep construction here so the cache cannot become stale.
supportGraphOf ::
    LA.Matrix Double ->
    Graph
supportGraphOf matrix =
    fromAdjacency
        dim
        [ ((i, j), entry > 0)
        | (i, row) <- zip [0 ..] rows
        , (j, entry) <- zip [0 ..] row
        ]
  where
    rows = LA.toLists matrix
    dim = length rows

{- | Matrix multiplication as transition composition: @p '<>' q@ takes a @p@
step followed by a @q@ step. Exact products preserve row-stochasticity and
associativity; 'Double' results are neither revalidated nor exactly
associative.
-}
instance Semigroup (TransitionMatrix state) where
    (<>) ::
        TransitionMatrix state ->
        TransitionMatrix state ->
        TransitionMatrix state
    p <> q = unsafeTransitionMatrix (unTransitionMatrix p LA.<> unTransitionMatrix q)

{- | The identity matrix represents zero transitions and is the unit of the
transition-composition monoid.
-}
instance (FiniteState state) => Monoid (TransitionMatrix state) where
    mempty :: TransitionMatrix state
    mempty = unsafeTransitionMatrix (LA.ident (stateCardinalityInt @state))