1
{-# LANGUAGE BangPatterns, CPP, FlexibleContexts, Rank2Types #-}
2
{-# OPTIONS_GHC -fsimpl-tick-factor=200 #-}
3
-- |
4
-- Module    : Statistics.Function
5
-- Copyright : (c) 2009, 2010, 2011 Bryan O'Sullivan
6
-- License   : BSD3
7
--
8
-- Maintainer  : bos@serpentine.com
9
-- Stability   : experimental
10
-- Portability : portable
11
--
12
-- Useful functions.
13
14
module Statistics.Function
15
    (
16
    -- * Scanning
17
      minMax
18
    -- * Sorting
19
    , sort
20
    , gsort
21
    , sortBy
22
    , partialSort
23
    -- * Indexing
24
    , indexed
25
    , indices
26
    -- * Bit twiddling
27
    , nextHighestPowerOfTwo
28
    -- * Comparison
29
    , within
30
    -- * Arithmetic
31
    , square
32
    -- * Vectors
33
    , unsafeModify
34
    -- * Combinators
35
    , for
36
    , rfor
37
    ) where
38
39
#include "MachDeps.h"
40
41
import Control.Monad.ST (ST)
42
import Data.Bits ((.|.), shiftR)
43
import qualified Data.Vector.Algorithms.Intro as I
44
import qualified Data.Vector.Generic as G
45
import qualified Data.Vector.Unboxed as U
46
import qualified Data.Vector.Unboxed.Mutable as M
47
import Numeric.MathFunctions.Comparison (within)
48
49
-- | Sort a vector.
50
sort :: U.Vector Double -> U.Vector Double
51
sort = G.modify I.sort
52
{-# NOINLINE sort #-}
53
54
-- | Sort a vector.
55
gsort :: (Ord e, G.Vector v e) => v e -> v e
56
gsort = G.modify I.sort
57
{-# INLINE gsort #-}
58
59
-- | Sort a vector using a custom ordering.
60
sortBy :: (G.Vector v e) => I.Comparison e -> v e -> v e
61
sortBy f = G.modify $ I.sortBy f
62
{-# INLINE sortBy #-}
63
64
-- | Partially sort a vector, such that the least /k/ elements will be
65
-- at the front.
66
partialSort :: (G.Vector v e, Ord e) =>
67
               Int -- ^ The number /k/ of least elements.
68
            -> v e
69
            -> v e
70
partialSort k = G.modify (`I.partialSort` k)
71
{-# SPECIALIZE partialSort :: Int -> U.Vector Double -> U.Vector Double #-}
72
73
-- | Return the indices of a vector.
74
indices :: (G.Vector v a, G.Vector v Int) => v a -> v Int
75
indices a = G.enumFromTo 0 (G.length a - 1)
76
{-# INLINE indices #-}
77
78
-- | Zip a vector with its indices.
79
indexed :: (G.Vector v e, G.Vector v Int, G.Vector v (Int,e)) => v e -> v (Int,e)
80
indexed a = G.zip (indices a) a
81
{-# INLINE indexed #-}
82
83
data MM = MM {-# UNPACK #-} !Double {-# UNPACK #-} !Double
84
85
-- | Compute the minimum and maximum of a vector in one pass.
86
minMax :: (G.Vector v Double) => v Double -> (Double, Double)
87
minMax = fini . G.foldl' go (MM (1/0) (-1/0))
88
  where
89
    go (MM lo hi) k = MM (min lo k) (max hi k)
90
    fini (MM lo hi) = (lo, hi)
91
{-# INLINE minMax #-}
92
93
-- | Efficiently compute the next highest power of two for a
94
-- non-negative integer.  If the given value is already a power of
95
-- two, it is returned unchanged.  If negative, zero is returned.
96
nextHighestPowerOfTwo :: Int -> Int
97
nextHighestPowerOfTwo n
98
#if WORD_SIZE_IN_BITS == 64
99
  = 1 + _i32
100
#else
101
  = 1 + i16
102
#endif
103
  where
104
    i0   = n - 1
105
    i1   = i0  .|. i0  `shiftR` 1
106
    i2   = i1  .|. i1  `shiftR` 2
107
    i4   = i2  .|. i2  `shiftR` 4
108
    i8   = i4  .|. i4  `shiftR` 8
109
    i16  = i8  .|. i8  `shiftR` 16
110
    _i32 = i16 .|. i16 `shiftR` 32
111
-- It could be implemented as
112
--
113
-- > nextHighestPowerOfTwo n = 1 + foldl' go (n-1) [1, 2, 4, 8, 16, 32]
114
--     where go m i = m .|. m `shiftR` i
115
--
116
-- But GHC do not inline foldl (probably because it's recursive) and
117
-- as result function walks list of boxed ints. Hand rolled version
118
-- uses unboxed arithmetic.
119
120
-- | Multiply a number by itself.
121
square :: Double -> Double
122
square x = x * x
123
124
-- | Simple for loop.  Counts from /start/ to /end/-1.
125
for :: Monad m => Int -> Int -> (Int -> m ()) -> m ()
126
for n0 !n f = loop n0
127
  where
128
    loop i | i == n    = return ()
129
           | otherwise = f i >> loop (i+1)
130
{-# INLINE for #-}
131
132
-- | Simple reverse-for loop.  Counts from /start/-1 to /end/ (which
133
-- must be less than /start/).
134
rfor :: Monad m => Int -> Int -> (Int -> m ()) -> m ()
135
rfor n0 !n f = loop n0
136
  where
137
    loop i | i == n    = return ()
138
           | otherwise = let i' = i-1 in f i' >> loop i'
139
{-# INLINE rfor #-}
140
141
unsafeModify :: M.MVector s Double -> Int -> (Double -> Double) -> ST s ()
142
unsafeModify v i f = do
143
  k <- M.unsafeRead v i
144
  M.unsafeWrite v i (f k)
145
{-# INLINE unsafeModify #-}