Haskell repa --- mapping with indices

Question

Imagine I want to map a function over an array, but the function has a type not just of a -> b but a -> Int -> b i.e. the function also takes an index. How do I do that?

Answer 1

Use zipWith

zipWith (\idx ele -> if even idx then div ele 2 else ele) [0..] xs

Answer 2

Good question, and it wasn't documented in the Repa tutorial, so I've updated it with a new section on traversals.

In particular, traverse lets you:

• change the shape of the output array
• index any eleemnt
• observe the current element

Meaning you can do things like:

Replace all eleemnts with their row index

> traverse a id (\_ (Z :. i :. j :. k) -> i) 
[0,0,0,0,0,0,0,0,0
,1,1,1,1,1,1,1,1,1
,2,2,2,2,2,2,2,2,2]

Multiply an element by its row

> traverse a id (\f (Z :. i :. j :. k) -> f (Z :. i :. j :. k) * i) 
[0,0,0,0,0,0,0,0,0
,10,11,12,13,14,15,16,17,18
,38,40,42,44,46,48,50,52,54]

And so on. travese is very powerful, and is also magically parallel.

Advanced: parallel image desaturation

Example from the Repa tutorial

Answer 3

Short answer, use traverse.

Longer example:

import qualified Data.Array.Repa as A
import qualified Data.Vector.Unboxed as U

arr1 :: A.Array A.DIM2 Double
arr1 = A.fromVector (A.Z A.:. 2 A.:. 3) $ U.fromList [1::Double,2,3,4,5,6]

arr2 :: A.Array A.DIM2 Double
arr2 = A.traverse arr1 id (\lf i@(A.Z A.:. r A.:. c) -> 
                  (lf i) + (fromIntegral r) + (fromIntegral c))  

arr1 is a 2x3 matrice. traverse is a function that takes (1) the original array, (2) a function for mapping source indices to target indices, and (3) a function that is given (i) a lookup function into the original array and (ii) an index that returns a new value.

So here arr2 modifies each of the original elements by adding the row and column indices of that particular entry.