DTW optimization with lazy evaluation

Dtw.hs

module Dtw where

import Data.Array

-- optmized using lazy evaluation, not using array
dtw :: Eq a => [a] -> [a] -> (a -> a -> Int) -> Maybe Int
dtw [] _ _ = Nothing
dtw _ [] _ = Nothing
dtw a b diff =
 let 
   -- diagonal from the top-left element
   mainDiag = oneDiag  a b diff (head uppers) ( 0 : (head lowers))

   -- diagonals above the mainDiag
   uppers   = eachDiag a b (mainDiag : uppers)

   -- diagonals below the mainDiag
   lowers   = eachDiag b a (mainDiag : lowers)  -- !

   eachDiag a [] diags = []
   eachDiag a (b:bs) (lastDiag:diags) =
    let nextDiag = head (tail diags)
    in  (oneDiag a bs diff nextDiag lastDiag):(eachDiag a bs diags)

   -- which is the diagonal containing the bottom R.H. elt?
   lab = (length a) - (length b)

 in Just $ last ( if lab == 0 
                  then mainDiag
                  else if lab > 0 
                       then lowers !! ( lab-1)
                       else uppers !! (-lab-1) )

min3 x y z =
  -- see L. Allison, Lazy Dynamic-Programming can be Eager
  --     Inf. Proc. Letters 43(4) pp207-212, Sept' 1992
  if x < y then x else min y z   -- makes it O(|a|*D(a,b))
  -- min x (min y z)             -- makes it O(|a|*|b|)
  -- the fast one does not always evaluate all three values.


-- compute one diagonal 
oneDiag :: Eq a => [a] -> [a] -> (a -> a -> Int) -> [Int] -> [Int] -> [Int]
oneDiag a b diff diagAbove diagBelow =  -- \   \   \
  let                                   --  \   \   \
    doDiag [] b nw n w = []             --   \  nw   n
    doDiag a [] nw n w = []             --    \   \
    doDiag (a:as) (b:bs) nw n w =       --      w   me
      let me = (diff a b) + (min3 (head w) nw (head n))
      in  me : (doDiag as bs me (tail n) (tail w))

    firstelt :: Int
    firstelt = (diff (head a) (head b))+(head diagBelow)
    
    thisdiag =
      firstelt:(doDiag a b firstelt diagAbove (tail diagBelow))

  in thisdiag


as :: [Int]
as = [0,0,0,3,6,13,25,22,7,2,1,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0]
bs :: [Int]
bs = [0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,4,5,12,24,23,8,3,1,0,0,0,0,0]


{-
$ ghci Dtw.hs

Ok, modules loaded: Dtw.
*Dtw> dtw as bs (\x y -> abs $ x - y)
21

-}

-- standard version using arrays

-- strict version
dtwS :: Eq a => [a] -> [a] -> (a -> a -> Int) -> Maybe Int
dtwS [] _ diff = Nothing
dtwS _ [] diff = Nothing
dtwS xs ys diff = Just $ table ! (m,n)      
  where 
    (m,n) = (length xs, length ys)
    x     = array (1,m) (zip [1..] xs)
    y     = array (1,n) (zip [1..] ys)
    
    table :: Array (Int,Int) Int
    table = array bnds [(ij, dist ij) | ij <- range bnds]
    bnds  = ((0,0),(m,n))

    dist (0,j) = 0
    dist (i,0) = 0
    dist (i,j) = (diff (x!i) (y!j)) + (minimum [table ! (i-1,j), table ! (i,j-1) , table ! (i-1,j-1)])