# statistics / Statistics / Function.hs

 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 {-# LANGUAGE Rank2Types, TypeOperators #-} -- | -- Module : Statistics.Function -- Copyright : (c) 2009 Bryan O'Sullivan -- License : BSD3 -- -- Maintainer : bos@serpentine.com -- Stability : experimental -- Portability : portable -- -- Useful functions. module Statistics.Function ( minMax , sort , partialSort -- * Array setup , createU , createIO ) where import Control.Exception (assert) import Control.Monad.ST (ST, unsafeIOToST, unsafeSTToIO) import Data.Array.Vector.Algorithms.Combinators (apply) import Data.Array.Vector import qualified Data.Array.Vector.Algorithms.Intro as I -- | Sort an array. sort :: (UA e, Ord e) => UArr e -> UArr e sort = apply I.sort {-# INLINE sort #-} -- | Partially sort an array, such that the least /k/ elements will be -- at the front. partialSort :: (UA e, Ord e) => Int -- ^ The number /k/ of least elements. -> UArr e -> UArr e partialSort k = apply (\a -> I.partialSort a k) {-# INLINE partialSort #-} data MM = MM {-# UNPACK #-} !Double {-# UNPACK #-} !Double -- | Compute the minimum and maximum of an array in one pass. minMax :: UArr Double -> Double :*: Double minMax = fini . foldlU 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 an array, using the given 'ST' action to populate each -- element. createU :: (UA e) => forall s. Int -> (Int -> ST s e) -> ST s (UArr e) createU size itemAt = assert (size >= 0) \$ newMU size >>= loop 0 where loop k arr | k >= size = unsafeFreezeAllMU arr | otherwise = do r <- itemAt k writeMU arr k r loop (k+1) arr {-# INLINE createU #-} -- | Create an array, using the given 'IO' action to populate each -- element. createIO :: (UA e) => Int -> (Int -> IO e) -> IO (UArr e) createIO size itemAt = unsafeSTToIO \$ createU size (unsafeIOToST . itemAt) {-# INLINE createIO #-}