Boundingcap.hs

module Boundingcap
       ( computeCap,
         theoreticCap2
       ) where

import qualified Data.Vector.Generic   as V
import           Debug.Trace
import           Numeric.LinearAlgebra as LA

for = flip map

row :: Matrix Double -> Int -> Vector Double
m `row` k = flatten $ m ? [k]

col :: Matrix Double -> Int -> Vector Double
m `col` j = flatten $ m ¿ [j]

computeCap :: Double -> Matrix Double -> (Double, Vector Double)
computeCap epsilon q = computeCap' epsilon q $ V.replicate n (1/fromIntegral n) where n = cols q

computeCap' :: Double -> Matrix Double -> Vector Double -> (Double, Vector Double)
computeCap' epsilon q p = (cap, p')
    where
        n = cols q
        m = rows q
        qp = q #> p
        c = V.generate n $ \j -> exp . sum . for [0..m-1] $ \k ->
                let qkj = q ! k ! j in
                qkj * logBase 2 (qkj/(qp ! k))
        il = log $ p `dot` c
        iu = log $ maxElement c
        (cap,p') = if iu - il < epsilon
            then (il,p)
            else
                computeCap' epsilon q newp where
                    pc = p `dot` c
                    newp = V.zipWith (((.).(.)) (/pc) (*)) p c

theoreticCap2 :: Matrix Double -> (Double, [Double])
theoreticCap2 q = (cap, [p0, 1-p0])
    where
        eps0 = q ! 1 ! 0
        eps1 = q ! 0 ! 1
        h2 eps = (-eps * logBase 2 eps) - ((1-eps) * logBase 2 (1-eps))
        z = 2**((h2 eps0 - h2 eps1)/(1-eps0-eps1))
        cap = eps0 / (1-eps0-eps1) * h2 eps1 - (1-eps1) / (1-eps0-eps1) * h2 eps0 + logBase 2 (1 + z)
        p0 = (1 - eps1 * (1+z))/((1-eps0-eps1)*(1+z))