Implementing 2D Torus Array with Haskell's Repa

Question

I just started looking at Repa and would like to know how to best implement a wrap-around, torus style 2D array read/written by stencil operation. I implemented this before using the ST monad and mutable vectors, but it seems this is not supported by Repa. A few approaches seem possible:

• I could just 'traverse' the array and do the index wrapping myself at each element. For simple stencils the cost of applying the wrapping logic everywhere is quite severe, so I'd need to avoid this one

• Data.Array.Repa.Stencil doesn't support the boundary condition I need, but looks like Data.Array.Repa.Algorithms.Convolve does. Documentation suggests a heavy performance penalty, though

• From what I understand, I can traverse over a subset of the array using slices. So, I could define an interior (1, 1) - (w-2, h-2), two horizontal and two vertical slabs representing the boundary, and then traverse them using two different functions / stencils, resulting in a single result array? Any sample code or further documentation on this?

• Repa also seems to have the concept of 'partitions', can they be used for implementing boundary conditions?

Which one is likely to be fastest? Any options I'm missing?

Thanks!

Answer 1

The most efficient way is to employ Partitioned array representation (your 4th option), however, it is inconvenient, because you should work with 5 areas by hand.

In Yarr you could write an utility

dim2WrapAround :: USource r l Dim2 a => UArray r l Dim2 a -> Dim2 -> Dim2 -> IO a
{-# INLINE dim2WrapAround #-}
dim2WrapAround arr (sizeX, sizeY) (posX, posY) =
    index arr (wrap sizeX posX, wrap sizeY posY)
  where wrap size pos = (pos + size) `mod` size

-- I'm afraid to write the signature...
{-# INLINE convolveOnThorus #-}
convolveOnThorus = convolveLinearDim2WithStaticStencil dim2WrapAround

Usage:

myConvolution :: UArray F L Dim2 Float -> UArray CV CVL Dim2 Float
myConvolution = convolveOnThorus [dim2St| some
                                           coeffs
                                               here |]