arrayshaskellrepa

Haskell repa --- mapping with indices

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?

Solution

  • 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.