Haskell Repa Package. Some Traverse Examples

This post provides some simple examples of the traverse function from the repa package. I found this function to look quite tricky and even frightening for a newbie so I share my experience with it. The repa package itself is so unbelievably awesome that I won’t tell anything about it, if you are interested check Repa tutorial.

We will need some dummy arrays to illustrate the usage of traverse. They will be provided by genArr function:

import Data.Array.Repa as R

genArr :: Int -> Array DIM1 Int
genArr n = fromList (Z :. n) [1..n]


*Main> genArr 12

As usual, the best point to start studying a function is a type of that function. From repa sources:

-- Generic Traversal -----------------------------------------------------------------------------
-- | Unstructured traversal.
    :: forall sh sh' a b
    .  (Shape sh, Shape sh', Elt a)
    => Array sh a                     -- ^ Source array.
    -> (sh  -> sh')                   -- ^ Function to produce the extent of the result.
    -> ((sh -> a) -> sh' -> b)        -- ^ Function to produce elements of the result. 
                                      --   It is passed a lookup function to get elements of the source.
    -> Array sh' b
{-# INLINE traverse #-}
traverse arr transExtent newElem
     = arr `deepSeqArray` 
          fromFunction (transExtent (extent arr)) (newElem (arr !))

It is readily seen that traverse requires three arguments. The first one is the array to be traversed. The second argument serves to specify the shape of the output array (from the shape of the input array). If these two shapes are the same, then the second argument should be id. The third argument might be the most confusing, it is in charge of filling the output array.

Now we come with the first example — how to produce an identical array with traverse:

identicalArray arr = traverse arr id id 

A slightly less stupid example — how to get the first s elements of an array:

firstS s arr = traverse arr (const $ Z :. s) id
*Main> firstS 4 $ genArr 12

The next example provides the function that adds to each element of an array its index:

plusI arr = traverse arr id $ \f (Z :. i) -> f (Z :. i) + i 

You might be confused with the last argument. If you are puzzled (like I was) “what does that f mean?” the answer is indeed provided by documentation — “… function to get elements of the source”. That is it is just index arr function and you may read

\f (Z :. i) -> f (Z :. i) + i 

in the last example as

\arr ! (Z :. i) -> arr ! (Z :. i) + i 

If you are not familiar with \f(x) -> ... syntax it is equivalent to \f -> \x -> ... .

And finally some justification:

*Main> plusI $ genArr 12
This entry was posted in Haskell and tagged , , . Bookmark the permalink.

Leave a Reply

Fill in your details below or click an icon to log in:

WordPress.com Logo

You are commenting using your WordPress.com account. Log Out / Change )

Twitter picture

You are commenting using your Twitter account. Log Out / Change )

Facebook photo

You are commenting using your Facebook account. Log Out / Change )

Google+ photo

You are commenting using your Google+ account. Log Out / Change )

Connecting to %s