如何沿最里面的维度对 massiv 数组进行排序?

How to sort a massiv array along innermost dimension?

我有一个 Array r Ix2 a 使得 (Manifest r Ix2 a, Ord a)。我想在最里面的维度上对这个数组进行排序——也就是说,在内部对每一行进行排序,而不是跨行排序。根据 this,massiv 根本没有实现任何排序。我是否必须自己动手,或者我可以重新使用 Vectors 已经存在的东西(例如 vector-algorithms)?

当然,推出自己的排序并向 massiv 库提交 PR 会更好;) 但有一种方法可以退回到 vector-algorithms 包。我很好奇我如何有效地做到这一点,在这里,以及对每一行进行排序的自动并行化:

{-# 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

我确实将此添加到此 commit 中的 massiv 中的示例以及一个简单的 属性 测试。