Relaxation labelling in Haskell

relaxation-labelling.hs

-- Relaxation labelling in Haskell

-- Usage: Load it in ghci, type `step 2` for 2 iterations.

import           Data.List.Split (chunksOf)
import           Text.Printf (printf)
import qualified Data.MemoCombinators as Memo -- data-memocombinators
import           Data.MemoCombinators (memo3)


data Label = Edge | NotEdge deriving (Eq, Show, Ord, Enum)

_N = 5
n = [0 .. _N^2-1]
labels = [Edge .. NotEdge]

taskA =
  [ [0,   0,   0,   0, 0]
  , [0, 0.1, 0.1, 0.1, 0]
  , [0, 0.1, 0.9,   0, 0]
  , [0, 0.1,   0,   0, 0]
  , [0,   0,   0,   0, 0]
  ]
taskB =
  [ [0,   0,   0,   0, 0]
  , [0,   0,   0,   0, 0]
  , [0,   0,   1,   0, 0]
  , [0,   0,   0,   0, 0]
  , [0,   0,   0,   0, 0]
  ]
taskC =
  [ [0,   0,   0,   0, 0]
  , [0, 0.1, 0.1, 0.1, 0]
  , [0, 0.1,   1, 0.1, 0]
  , [0, 0.1, 0.1, 0.1, 0]
  , [0,   0,   0,   0, 0]
  ]
row x = x `quot` _N
col x = x `rem`  _N

-- What image to run on
image = taskC

-- r _            _            = 1 -- comment in for taskA
r (_, Edge)    (_, Edge)    = 2
r (_, Edge)    (_, NotEdge) = 1
r (_, NotEdge) (_, Edge)    = 1
r (_, NotEdge) (_, NotEdge) = 1


c i j
  | i == j    = 0
  | otherwise = if (abs (row i - row j) <= 1 &&
                    abs (col i - col j) <= 1   ) then 1 else 0

-- CONTEXTUAL SUPPORT q   for object ai having label λl
q s j ai λl =     sum      [  r (ai, λl) (j, λk)  *  p s j λk
            | λk <- labels ]

-- TOTAL SUPPORT Q   for object ai having label λl
qq s ai λl =   sum  [  c ai j  *  q s j ai λl
           | j <- n ]


-- Initial probabilities (step s = 0)
p' 0 ai Edge     = image !! row ai !! col ai
p' 0 ai NotEdge = 1 - p 0 ai Edge

-- Probability P   for object ai having label λl (step s)
p' s ai λl =            qq (s-1) ai λl  *  p (s-1) ai λl
                                        /
                      sum      [  qq (s-1) ai λk * p (s-1) ai λk
                | λk <- labels ]

-- Memoization to make it fast
p = memo3 enum enum enum p'

step s = display $ chunksOf _N [ p s i Edge | i <- n ]

display grid = mapM_ (putStrLn . unwords . map (' ':) . map format) grid
format = printf "%.5f" :: Double -> String



-- External code

-- Remove once https://github.com/luqui/data-memocombinators/issues/8
enum :: Enum a => Memo.Memo a
enum = Memo.wrap toEnum fromEnum Memo.integral