{-# LANGUAGE FlexibleContexts #-}
module ELynx.Data.MarkovProcess.RateMatrix
( RateMatrix,
ExchangeabilityMatrix,
StationaryDistribution,
isValid,
normalizeSD,
totalRate,
totalRateWith,
normalize,
normalizeWith,
setDiagonal,
toExchangeabilityMatrix,
fromExchangeabilityMatrix,
getStationaryDistribution,
exchFromListLower,
exchFromListUpper,
)
where
import qualified Data.Vector.Storable as V
import Numeric.LinearAlgebra hiding (normalize)
import Numeric.SpecFunctions
import Prelude hiding ((<>))
type RateMatrix = Matrix R
type ExchangeabilityMatrix = Matrix R
type StationaryDistribution = Vector R
epsRelaxed :: Double
epsRelaxed = 1e-5
isValid :: StationaryDistribution -> Bool
isValid d = epsRelaxed > abs (norm_1 d - 1.0)
normalizeSD :: StationaryDistribution -> StationaryDistribution
normalizeSD d = d / scalar (norm_1 d)
matrixSetDiagToZero :: Matrix R -> Matrix R
matrixSetDiagToZero m = m - diag (takeDiag m)
{-# INLINE matrixSetDiagToZero #-}
totalRateWith :: StationaryDistribution -> RateMatrix -> Double
totalRateWith d m = norm_1 $ d <# matrixSetDiagToZero m
totalRate :: RateMatrix -> Double
totalRate m = totalRateWith (getStationaryDistribution m) m
normalize :: RateMatrix -> RateMatrix
normalize m = normalizeWith (getStationaryDistribution m) m
normalizeWith :: StationaryDistribution -> RateMatrix -> RateMatrix
normalizeWith d m = scale (1.0 / totalRateWith d m) m
setDiagonal :: RateMatrix -> RateMatrix
setDiagonal m = diagZeroes - diag (fromList rowSums)
where
diagZeroes = matrixSetDiagToZero m
rowSums = map norm_1 $ toRows diagZeroes
toExchangeabilityMatrix ::
RateMatrix -> StationaryDistribution -> ExchangeabilityMatrix
toExchangeabilityMatrix m f = m <> diag oneOverF
where
oneOverF = cmap (1.0 /) f
fromExchangeabilityMatrix ::
ExchangeabilityMatrix -> StationaryDistribution -> RateMatrix
fromExchangeabilityMatrix em d = setDiagonal $ em <> diag d
eps :: Double
eps = 1e-12
normalizeSumVec :: V.Vector Double -> V.Vector Double
normalizeSumVec v = V.map (/ s) v
where
s = V.sum v
{-# INLINE normalizeSumVec #-}
getStationaryDistribution :: RateMatrix -> StationaryDistribution
getStationaryDistribution m =
if eps > abs (magnitude (eVals ! i))
then normalizeSumVec distReal
else
error
"getStationaryDistribution: Could not retrieve stationary distribution."
where
(eVals, eVecs) = eig (tr m)
i = minIndex eVals
distComplex = toColumns eVecs !! i
distReal = cmap realPart distComplex
ijToKLower :: Int -> Int -> Int
ijToKLower i j
| i > j = round (i `choose` 2) + j
| otherwise = error "ijToKLower: not defined for upper triangular matrix."
ijToKUpper :: Int -> Int -> Int -> Int
ijToKUpper n i j
| i < j = i * (n - 2) - round (i `choose` 2) + j - 1
| otherwise = error "ijToKUpper: not defined for lower triangular matrix."
fromListBuilderLower :: RealFrac a => [a] -> a -> a -> a
fromListBuilderLower es i j
| i > j = es !! ijToKLower iI jI
| i == j = 0.0
| i < j = es !! ijToKLower jI iI
| otherwise =
error
"Float indices could not be compared during matrix creation."
where
iI = round i :: Int
jI = round j :: Int
fromListBuilderUpper :: RealFrac a => Int -> [a] -> a -> a -> a
fromListBuilderUpper n es i j
| i < j = es !! ijToKUpper n iI jI
| i == j = 0.0
| i > j = es !! ijToKUpper n jI iI
| otherwise =
error
"Float indices could not be compared during matrix creation."
where
iI = round i :: Int
jI = round j :: Int
checkEs :: RealFrac a => Int -> [a] -> [a]
checkEs n es
| length es == nExp = es
| otherwise = error eStr
where
nExp = round (n `choose` 2)
eStr =
unlines
[ "exchFromListlower: the number of exchangeabilities does not match the matrix size",
"matrix size: " ++ show n,
"expected number of exchangeabilities: " ++ show nExp,
"received number of exchangeabilities: " ++ show (length es)
]
exchFromListLower :: (RealFrac a, Container Vector a) => Int -> [a] -> Matrix a
exchFromListLower n es = build (n, n) (fromListBuilderLower (checkEs n es))
exchFromListUpper :: (RealFrac a, Container Vector a) => Int -> [a] -> Matrix a
exchFromListUpper n es = build (n, n) (fromListBuilderUpper n (checkEs n es))