diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,19 @@ +Copyright (c) 2016 Jared Tobin + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in +all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +THE SOFTWARE. diff --git a/Setup.hs b/Setup.hs new file mode 100644 --- /dev/null +++ b/Setup.hs @@ -0,0 +1,2 @@ +import Distribution.Simple +main = defaultMain diff --git a/bench/Main.hs b/bench/Main.hs new file mode 100644 --- /dev/null +++ b/bench/Main.hs @@ -0,0 +1,16 @@ +{-# OPTIONS_GHC -fno-warn-type-defaults #-} + +module Main where + +import Criterion.Main +import Numeric.Sampling (sampleIO, resampleIO) + +main :: IO () +main = defaultMain [ + env (return ([1..10000] :: [Int])) $ \x -> + bgroup "benchmarks" [ + bench "sample" $ nfIO (sampleIO 100 x) + , bench "resample" $ nfIO (resampleIO 100 x) + ] + ] + diff --git a/lib/Numeric/Sampling.hs b/lib/Numeric/Sampling.hs new file mode 100644 --- /dev/null +++ b/lib/Numeric/Sampling.hs @@ -0,0 +1,104 @@ +{-# OPTIONS_GHC -Wall #-} +{-# OPTIONS_GHC -fno-warn-type-defaults #-} +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE FlexibleContexts #-} + +module Numeric.Sampling ( + -- * Without replacement + sample + , sampleIO + + -- * With replacement + , resample + , resampleIO + + -- * Unequal probability, with replacement + , presample + , presampleIO + + -- * Re-exported + , module System.Random.MWC + ) where + +import qualified Control.Foldl as F +import Control.Monad.Primitive (PrimMonad, PrimState) +import qualified Data.Foldable as Foldable +#if __GLASGOW_HASKELL__ < 710 +import Data.Foldable (Foldable) +#endif +import Data.Function (on) +import Data.List (sortBy) +import Data.Vector (Vector) +import Numeric.Sampling.Internal +import System.Random.MWC + +-- | (/O(n)/) Sample uniformly, without replacement. +-- +-- Returns Nothing if the desired sample size is larger than the collection +-- being sampled from. +sample + :: (PrimMonad m, Foldable f) + => Int -> f a -> Gen (PrimState m) -> m (Maybe (Vector a)) +sample n xs gen + | n < 0 = return Nothing + | otherwise = F.foldM (randomN n gen) xs +{-# INLINABLE sample #-} + +-- | (/O(n)/) 'sample' specialized to IO. +sampleIO :: Foldable f => Int -> f a -> IO (Maybe (Vector a)) +sampleIO n xs = do + gen <- createSystemRandom + sample n xs gen +{-# INLINABLE sampleIO #-} + +-- | (/O(n log n)/) Sample uniformly with replacement (bootstrap). +resample + :: (PrimMonad m, Foldable f) + => Int -> f a -> Gen (PrimState m) -> m [a] +resample n xs = presample n weighted where + weight = recip (F.fold F.genericLength xs) + weighted = zip (repeat weight) (Foldable.toList xs) +{-# INLINABLE resample #-} + +-- | (/O(n log n)/) 'resample' specialized to IO. +resampleIO :: (Foldable f) => Int -> f a -> IO [a] +resampleIO n xs = do + gen <- createSystemRandom + resample n xs gen +{-# INLINABLE resampleIO #-} + +-- | (/O(n log n)/) Unequal probability resampling. +presample + :: (PrimMonad m, Foldable f) + => Int -> f (Double, a) -> Gen (PrimState m) -> m [a] +presample n weighted gen + | n <= 0 = return [] + | otherwise = do + let (bprobs, vals) = unzip $ sortProbs weighted + probs = drop 1 (F.scan F.sum bprobs) + cumulative = zip probs vals + computeSample n cumulative gen + where + computeSample + :: PrimMonad m => Int -> [(Double, a)] -> Gen (PrimState m) -> m [a] + computeSample size xs g = go [] size where + go !acc s + | s <= 0 = return acc + | otherwise = do + z <- uniform g + case F.fold (F.find ((>= z) . fst)) xs of + Just (_, val) -> go (val:acc) (pred s) + Nothing -> return acc + + sortProbs :: (Foldable f, Ord a) => f (a, b) -> [(a, b)] + sortProbs = sortBy (compare `on` fst) . Foldable.toList +{-# INLINABLE presample #-} + +-- | (/O(n log n)/) 'presample' specialized to IO. +presampleIO :: (Foldable f) => Int -> f (Double, a) -> IO [a] +presampleIO n weighted = do + gen <- createSystemRandom + presample n weighted gen +{-# INLINABLE presampleIO #-} + diff --git a/lib/Numeric/Sampling/Internal.hs b/lib/Numeric/Sampling/Internal.hs new file mode 100644 --- /dev/null +++ b/lib/Numeric/Sampling/Internal.hs @@ -0,0 +1,82 @@ +{-# LANGUAGE FlexibleContexts #-} + +-- Much of the code in this module is a modification of that found in the +-- 'foldl' library by Gabriel Gonzalez. Its license is reproduced below. + +-- Copyright (c) 2013 Gabriel Gonzalez +-- All rights reserved. +-- +-- Redistribution and use in source and binary forms, with or without modification, +-- are permitted provided that the following conditions are met: +-- * Redistributions of source code must retain the above copyright notice, +-- this list of conditions and the following disclaimer. +-- * Redistributions in binary form must reproduce the above copyright notice, +-- this list of conditions and the following disclaimer in the documentation +-- and/or other materials provided with the distribution. +-- * Neither the name of Gabriel Gonzalez nor the names of other contributors +-- may be used to endorse or promote products derived from this software +-- without specific prior written permission. +-- +-- THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND +-- ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED +-- WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +-- DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR +-- ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES +-- (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; +-- LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON +-- ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +-- (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +-- SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +module Numeric.Sampling.Internal ( + randomN + ) where + +import Control.Foldl (FoldM (..)) +import Control.Monad (when) +import Control.Monad.Primitive +import Data.Vector.Generic (Mutable, Vector) +import qualified Data.Vector.Generic as V +import Data.Vector.Generic.Mutable (MVector) +import qualified Data.Vector.Generic.Mutable as M +import System.Random.MWC + +data VectorState = Incomplete {-# UNPACK #-} !Int | Complete + +data RandomNState m v a = RandomNState + { _size :: !VectorState + , _reservoir :: !(Mutable v (PrimState m) a) + , _position :: {-# UNPACK #-} !Int + , _gen :: !(Gen (PrimState m)) + } + +randomN + :: (PrimMonad m, Vector v a) + => Int -> Gen (PrimState m) -> FoldM m a (Maybe (v a)) +randomN n gen = FoldM step begin done where + step + :: (PrimMonad m, MVector (Mutable v) a) + => RandomNState m v a -> a -> m (RandomNState m v a) + step (RandomNState (Incomplete m) mv i g) a = do + M.write mv m a + let m' = m + 1 + let s = if n <= m' then Complete else Incomplete m' + return $! RandomNState s mv (i + 1) g + + step (RandomNState Complete mv i g) a = do + r <- uniformR (0, i - 1) g + when (r < n) (M.unsafeWrite mv r a) + return (RandomNState Complete mv (i + 1) g) + + begin = do + mv <- M.new n + let s = if n <= 0 then Complete else Incomplete 0 + return (RandomNState s mv 1 gen) + + done :: (PrimMonad m, Vector v a) => RandomNState m v a -> m (Maybe (v a)) + done (RandomNState (Incomplete _) _ _ _) = return Nothing + done (RandomNState Complete mv _ _) = do + v <- V.freeze mv + return (Just v) +{-# INLINABLE randomN #-} + diff --git a/sampling.cabal b/sampling.cabal new file mode 100644 --- /dev/null +++ b/sampling.cabal @@ -0,0 +1,63 @@ +name: sampling +version: 0.1.1 +synopsis: Sample values from collections. +homepage: https://github.com/jtobin/sampling +license: MIT +license-file: LICENSE +author: Jared Tobin +maintainer: jared@jtobin.ca +category: Math +build-type: Simple +cabal-version: >=1.10 +description: + Basic sampling tools. + . + Exports variations on two simple functions for sampling from arbitrary + 'Foldable' collections: + . + * 'sample', for sampling without replacement + . + * 'resample', for sampling with replacement (i.e., a bootstrap) + +Source-repository head + Type: git + Location: http://github.com/jtobin/sampling.git + +library + default-language: Haskell2010 + hs-source-dirs: lib + ghc-options: + -Wall + other-modules: + Numeric.Sampling.Internal + exposed-modules: + Numeric.Sampling + build-depends: + base < 5 + , foldl + , mwc-random + , primitive + , vector + +executable sampling-test + hs-source-dirs: src + Main-is: Main.hs + default-language: Haskell2010 + ghc-options: + -Wall -O2 + build-depends: + base + , sampling + +benchmark bench-sampling + type: exitcode-stdio-1.0 + hs-source-dirs: bench + Main-is: Main.hs + default-language: Haskell2010 + ghc-options: + -Wall -O2 + build-depends: + base + , criterion + , sampling + diff --git a/src/Main.hs b/src/Main.hs new file mode 100644 --- /dev/null +++ b/src/Main.hs @@ -0,0 +1,10 @@ +{-# OPTIONS_GHC -fno-warn-type-defaults #-} + +module Main where + +import Numeric.Sampling + +main :: IO () +main = do + foo <- resampleIO 100 ([1..100000] :: [Int]) + print foo