How to sort a massiv array along innermost dimension?

Question

I have an Array r Ix2 a such that (Manifest r Ix2 a, Ord a). I'd like to sort this array on its innermost dimension - that is, sort each row internally, but not across rows. According to this, massiv doesn't have any sorting implemented at all. Would I have to roll my own, or can I re-use something that already exists for Vectors (such as vector-algorithms for example)?

Answer 1

Of course, it would be better to roll out your own sorting and submit a PR to massiv library ;) But there is a way to fall back onto vector-algorithms package. I was curious how I'd do it efficiently and here it is, along with automatic parallelization of sorting each row:

{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies        #-}
module Examples.SortRows where

import           Data.Massiv.Array                 as A
import           Data.Massiv.Array.Manifest.Vector as A
import           Data.Massiv.Core.Scheduler
import           Data.Typeable
import           Data.Vector.Algorithms.Merge
import           Data.Vector.Generic               as VG
import           Data.Vector.Generic.Mutable       as VGM
import           System.IO.Unsafe

sortRows ::
     forall r e v.
     (Ord e, Typeable v, A.Mutable r Ix2 e, VG.Vector v e, ARepr v ~ r, VRepr r ~ v)
  => Array r Ix2 e
  -> Array r Ix2 e
sortRows arr = unsafePerformIO $ do
  mv :: VG.Mutable v RealWorld e <- VG.thaw (A.toVector arr :: v e)
  let comp = getComp arr
      sz@(m :. n) = size arr
  case comp of
    Seq -> do
      loopM_ 0 (< m) (+ 1) $ \i -> sort $ VGM.slice (toLinearIndex sz (i :. 0)) n mv
    ParOn wIds ->
      withScheduler_ wIds $ \scheduler -> do
        loopM_ 0 (< m) (+ 1) $ \i ->
          scheduleWork scheduler $ sort $ VGM.slice (toLinearIndex sz (i :. 0)) n mv
  v :: v e <- VG.unsafeFreeze mv
  return $ A.fromVector comp sz v

I did add this to examples in massiv in this commit together with a simple property test.