Haskell data analysis: Markov matrix calculation

calc-markov-matrices.hs

module Main where

import Control.Applicative ((<$>))
import Control.Monad (forM_, zipWithM_)
import Data.List (intercalate)
import System.FilePath
import qualified Data.Vector.Storable as SV

import Data.NetCDF
import Data.NetCDF.HMatrix
import Foreign.C
import Numeric.LinearAlgebra.HMatrix

import Debug.Trace

type HMatrixRet a = IO (Either NcError (HMatrix a))
type M = Matrix Double

basedir, workdir :: FilePath
basedir = "/big/project-storage/haskell-data-analysis/atmos-non-diffusive"
workdir = basedir </> "dynamics"

partitionSizes :: [Int]
partitionSizes = [4, 6, 4, 4]

main :: IO ()
main = do
  -- Open partition time series data file for input.
  Right innc <- openFile $ workdir </> "partitions.nc"
  let Just ntime = ncDimLength <$> ncDim innc "time"
      Just npart = ncDimLength <$> ncDim innc "part"
  let (Just partvar) = ncVar innc "partition"
  Right (HMatrix partin) <- get innc partvar :: HMatrixRet CInt

  -- Calculate Markov matrix for each partition.
  forM_ (zip3 (toColumns partin) partitionSizes [1..]) $ \(part, ps, ip) -> do
    putStrLn $ "PARTITION " ++ show ip ++ "\n"
    let (mT, iNk) = transMatrix ps $ SV.map fromIntegral part
        mM = cmap (/ (fromIntegral iNk)) mT
        (mTA, mTS) = splitMarkovMatrix mT
        (mMA, mMS) = splitMarkovMatrix mM
        mS = stdDevMatrix iNk mM
        sig = sigLevels mTS mTA mS
    putStrLn "T:"
    disp 3 mT
    putStrLn "\nM^A + |M^A| (x 100):"
    disp 3 $ scale 100.0 $ mMA + abs mMA
    putStrLn "\nSignificance levels:"
    disp 3 $ sig
    putStrLn "\n--------------------------------------------------"
    latexOutSimple "\\mathbf{T}" mT
    putStrLn ""
    latexOutSig "\\mathbf{M}^A + |\\mathbf{M}^A|" "\\frac{1}{100}"
      (scale 100.0 $ mMA + abs mMA) sig
    putStrLn "--------------------------------------------------\n"


latexOutSimple :: String -> M -> IO ()
latexOutSimple lab m = do
  putStrLn $ lab ++ " = \\begin{pmatrix}"
  forM_ (toRows m) $ \r ->
    putStrLn $ "  " ++ (intercalate " & " $ map show $ SV.toList r) ++ " \\\\"
  putStrLn "\\end{pmatrix}"

latexOutSig :: String -> String -> M -> M -> IO ()
latexOutSig lab1 lab2 m sig = do
  putStrLn $ lab1 ++ " = " ++ lab2 ++ " \\begin{pmatrix}"
  zipWithM_ dorow (toRows m) (toRows sig)
  putStrLn "\\end{pmatrix}"
  where dorow r s =
          let xs = intercalate " & " $ zipWith doone (SV.toList r) (SV.toList s)
          in putStrLn $ "  " ++ xs ++ " \\\\"
        doone x s = let sx = show (fromIntegral (round $ x * 10) / 10 :: Double)
                    in if sx == "0.0"
                       then "0"
                       else addsig s sx
        addsig s sx
          | s == 95.0 = "\\underline{\\underline{\\mathbf{" ++ sx ++ "}}}"
          | s == 90.0 = "\\underline{\\mathbf{" ++ sx ++ "}}"
          | s == 85.0 = "\\mathbf{" ++ sx ++ "}"
          | s == 80.0 = "\\underline{" ++ sx ++ "}"
          | s == 75.0 = sx
          | otherwise = "\\mathit{" ++ sx ++ "}"

transMatrix :: Int -> SV.Vector Int -> (M, Int)
transMatrix n pm = (accum (konst 0.0 (n, n)) (+) $ zip steps (repeat 1.0), ns)
  where allSteps = [((pm SV.! (i + 1)) - 1, (pm SV.! i) - 1) |
                    i <- [0..SV.length pm - 2], (i + 1) `mod` 21 /= 0]
        steps0 = map (\k -> filter (\(i, j) -> i == k) allSteps) [0..n-1]
        ns = minimum $ map length steps0
        steps = concat $ map (take ns) steps0

splitMarkovMatrix :: M -> (M, M)
splitMarkovMatrix mm = (a, s)
  where s = scale 0.5 $ mm + tr mm
        a = scale 0.5 $ mm - tr mm

stdDevMatrix :: Int -> M -> M
stdDevMatrix iNk mMS = cmap f mMS
  where f x = sqrt (x * (1.0 - x) * fromIntegral iNk)

sigLevels :: M -> M -> M -> M
sigLevels mTS mTA mS =
  reshape (cols mTS) $ SV.zipWith3 f (flatten mTS) (flatten mTA) (flatten mS)
  where f xTS xTA xS = if xTA > 0 && xTS >= 10
                       then lookupSig (abs xTA) xS
                       else 0.0
        lookupSig x s
          | x > 1.96 * s = 95.0
          | x > 1.64 * s = 90.0
          | x > 1.44 * s = 85.0
          | x > 1.28 * s = 80.0
          | x > 1.15 * s = 75.0
          | otherwise    = 0.0

partition-phase-space.hs

module Main where

import Control.Applicative ((<$>))
import Control.Monad (forM_)
import System.FilePath
import qualified Data.Map as M

import Data.NetCDF
import Data.NetCDF.HMatrix
import Foreign.C
import qualified Data.NetCDF.Vector as V
import qualified Data.Vector.Storable as SV
import Numeric.LinearAlgebra.HMatrix


type SVRet a = IO (Either NcError (SV.Vector a))
type HMatrixRet a = IO (Either NcError (HMatrix a))

basedir, indir, outdir :: FilePath
basedir = "/big/project-storage/haskell-data-analysis/atmos-non-diffusive"
indir = basedir </> "spherical-pdf-decorrelated"
outdir = basedir </> "dynamics"

-- Partition component: theta range, phi range.
data Component = C { thmin :: Double, thmax :: Double
                   , phmin :: Double, phmax :: Double
                   } deriving Show

-- A partition is a list of components that cover the unit sphere.
type Partition = [Component]

-- Angle for 1-4-1 partition.
th4 :: Double
th4 = acos $ 2.0 / 3.0

-- Partitions.
partitions :: [Partition]
partitions = [ [ C 0        (pi/3)   0  (2*pi)
               , C (pi/3)   (2*pi/3) 0  pi
               , C (pi/3)   (2*pi/3) pi (2*pi)
               , C (2*pi/3) pi       0  (2*pi) ]
             , [ C 0        th4      0        (2*pi)
               , C th4      (pi-th4) 0        (pi/2)
               , C th4      (pi-th4) (pi/2)   pi
               , C th4      (pi-th4) pi       (3*pi/2)
               , C th4      (pi-th4) (3*pi/2) (2*pi)
               , C (pi-th4) pi       0        (2*pi) ]
             , [ C 0      (pi/2) 0  pi
               , C 0      (pi/2) pi (2*pi)
               , C (pi/2) pi     0  pi
               , C (pi/2) pi     pi (2*pi) ]
             , [ C 0      (pi/2) (pi/4)   (5*pi/4)
               , C 0      (pi/2) (5*pi/4) (pi/4)
               , C (pi/2) pi     (pi/4)   (5*pi/4)
               , C (pi/2) pi     (5*pi/4) (pi/4) ] ]

npartitions :: Int
npartitions = length partitions

-- Convert list of (theta, phi) coordinates to partition component
-- numbers for a given partition.
convert :: Partition -> [(Double, Double)] -> [Int]
convert part pts = map (convOne part) pts
  where convOne comps (th, ph) = 1 + length (takeWhile not $ map isin comps)
          where isin (C thmin thmax ph1 ph2) =
                  if ph1 < ph2
                  then th >= thmin && th < thmax && ph >= ph1 && ph < ph2
                  else th >= thmin && th < thmax && (ph >= ph1 || ph < ph2)

main :: IO ()
main = do
  -- Open projected points data file for input.
  Right innc <- openFile $ indir </> "sphere-proj.nc"
  let Just ntime = ncDimLength <$> ncDim innc "time"
  let (Just projvar) = ncVar innc "proj"
      (Just timevar) = ncVar innc "time"
  Right time <- get innc timevar :: SVRet CDouble
  Right (HMatrix projsin) <-
    getA innc projvar [0, 0] [ntime, 2] :: HMatrixRet CDouble

  -- Convert (theta, phi) coordinates to partition component indexes
  -- for each partition.
  let phth = toColumns $ cmap realToFrac projsin :: [Vector Double]
      pts = zip (SV.toList $ phth !! 0) (map (+pi) $ SV.toList $ phth !! 1)
      idxs = flatten $ fromColumns $
             map (SV.fromList . map fromIntegral) $
             map (flip convert pts) partitions :: SV.Vector CInt

  -- Set up output file.
  let outpartdim = NcDim "part" npartitions False
      outpartvar = NcVar "part" NcInt [outpartdim] M.empty
      outtimedim = NcDim "time" 0 True
      outtimevar = NcVar "time" NcDouble [outtimedim] (ncVarAttrs timevar)
      outvar = NcVar "partition" NcInt [outtimedim, outpartdim] M.empty
      outncinfo =
        emptyNcInfo (outdir </> "partitions.nc") #
        addNcDim outtimedim # addNcDim outpartdim #
        addNcVar outtimevar # addNcVar outpartvar #
        addNcVar outvar

  flip (withCreateFile outncinfo) (putStrLn . ("ERROR: " ++) . show) $
    \outnc -> do
      -- Write coordinate variable values.
      putA outnc outtimevar [0] [ntime] time
      put outnc outpartvar $ (SV.enumFromN 1 npartitions :: SV.Vector CInt)
      put outnc outvar idxs
      return ()