Haskell repa - how to reduce array and return index?

Question

In GNU Octave this code -

[e, ix] = min(X);

will return minimum element and it's location. How do you this in repa for arbitrary binary function?

This is what I came up with:

min x = z $ foldl' f (e,0,0) es
  where
    (e:es) = toList x
    f (a,ix,r) b = let ix' = ix+1 in if a < b then (a,ix',r) else (b,ix',ix')
    z (a,ix,r) = (a,r)

In above example we convert repa 1D matrix to list and use foldl' (from Data.List) with two accumulators - one for counting iterations (ix) and other to save position of min element (r). But the whole point of using repa is to use arrays, not lists!

In repa there are two folds for Array type (foldS and foldP) - but they can only take function of type (a -> a -> a) - meaning, I cannot pass tuple with accumulators to it. There is also traverse, which can, in principle, reduce 1D array to a scalar array:

min x = traverse x to0D min
  where
    to0D (Z:.i) = Z
    min f (Z) = ??? -- how to get elements for comparison?

The first thing that comes to mind is

[f (Z:.i) | i <- [1..n]], where n = (\(Z:.i) -> i) $ extent x

But this will also convert array to list, instead of doing computation on the array.

Answer 1

I'm no expert on Repa, but this seems to work for 1-D arrays. It can probably be adapted for other dimensions.

import Data.Array.Repa

indexed arr = traverse arr id (\src idx@(Z :. i) -> (src idx, i))

minimize arr = foldP f h t
  where
    (Z :. n) = extent arr
    arr' = indexed arr
    h = arr' ! (Z :. 0)
    t = extract (Z :. 1) (Z :. (n-1)) arr'
    f min@(valMin, iMin) x@(val, i) | val < valMin = x
                                    | otherwise = min