1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
|
{-# LANGUAGE Rank2Types, TypeOperators #-}
-- |
-- Module : Statistics.Function
-- Copyright : (c) 2009, 2010 Bryan O'Sullivan
-- License : BSD3
--
-- Maintainer : bos@serpentine.com
-- Stability : experimental
-- Portability : portable
--
-- Useful functions.
module Statistics.Function
(
minMax
, sort
, partialSort
, indexed
, indices
-- * Vector setup
, create
) where
import Control.Exception (assert)
import Control.Monad.Primitive (PrimMonad)
import Data.Vector.Algorithms.Combinators (apply)
import Data.Vector.Generic (unsafeFreeze)
import qualified Data.Vector.Algorithms.Intro as I
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Unboxed.Mutable as MU
-- | Sort a vector.
sort :: (U.Unbox e, Ord e) => U.Vector e -> U.Vector e
sort = apply I.sort
{-# INLINE sort #-}
-- | Partially sort a vector, such that the least /k/ elements will be
-- at the front.
partialSort :: (U.Unbox e, Ord e) =>
Int -- ^ The number /k/ of least elements.
-> U.Vector e
-> U.Vector e
partialSort k = apply (\a -> I.partialSort a k)
{-# INLINE partialSort #-}
-- | Return the indices of a vector.
indices :: (U.Unbox a) => U.Vector a -> U.Vector Int
indices a = U.enumFromTo 0 (U.length a - 1)
{-# INLINE indices #-}
-- | Zip a vector with its indices.
indexed :: U.Unbox e => U.Vector e -> U.Vector (Int,e)
indexed a = U.zip (indices a) a
{-# INLINE indexed #-}
data MM = MM {-# UNPACK #-} !Double {-# UNPACK #-} !Double
-- | Compute the minimum and maximum of a vector in one pass.
minMax :: U.Vector Double -> (Double , Double)
minMax = fini . U.foldl go (MM (1/0) (-1/0))
where
go (MM lo hi) k = MM (min lo k) (max hi k)
fini (MM lo hi) = (lo, hi)
{-# INLINE minMax #-}
-- | Create a vector, using the given action to populate each
-- element.
create :: (PrimMonad m, U.Unbox e) => Int -> (Int -> m e) -> m (U.Vector e)
create size itemAt = assert (size >= 0) $
MU.new size >>= loop 0
where
loop k arr | k >= size = unsafeFreeze arr
| otherwise = do r <- itemAt k
MU.write arr k r
loop (k+1) arr
{-# INLINE create #-}
|
