<commit_before>{-# LANGUAGE BangPatterns #-}
-- |
-- Module    : System.Random.MWC.Distributions
-- Copyright : (c) 2012 <NAME>
-- License   : BSD3
--
-- Maintainer  : <EMAIL>
-- Stability   : experimental
-- Portability : portable
--
-- Pseudo-random number generation for non-uniform distributions.

module System.Random.MWC.Distributions
    (
    -- * Variates: non-uniformly distributed values
      normal
    , standard
    , exponential
    , truncatedExp
    , gamma
    , chiSquare
    , geometric0
    , geometric1

    -- * References
    -- $references
    ) where

import Control.Monad (liftM)
import Control.Monad.Primitive (PrimMonad, PrimState)
import Data.Bits ((.&.))
import Data.Word (Word32)
import System.Random.MWC (Gen, uniform)
import qualified Data.Vector.Unboxed as I

-- Unboxed 2-tuple
data T = T {-# UNPACK #-} !Double {-# UNPACK #-} !Double


-- | Generate a normally distributed random variate with given mean
-- and standard deviation.
normal :: PrimMonad m
       => Double                -- ^ Mean
       -> Double                -- ^ Standard deviation
       -> Gen (PrimState m)
       -> m Double
{-# INLINE normal #-}
normal m s gen = do
  x <- standard gen
  return $! m + s * x

-- | Generate a normally distributed random variate with zero mean and
-- unit variance.
--
-- The implementation uses Doornik's modified ziggurat algorithm.
-- Compared to the ziggurat algorithm usually used, this is slower,
-- but generates more independent variates that pass stringent tests
-- of randomness.
standard :: PrimMonad m => Gen (PrimState m) -> m Double
{-# INLINE standard #-}
standard gen = loop
  where
    loop = do
      u  <- (subtract 1 . (*2)) `liftM` uniform gen
      ri <- uniform gen
      let i  = fromIntegral ((ri :: Word32) .&. 127)
          bi = I.unsafeIndex blocks i
          bj = I.unsafeIndex blocks (i+1)
      case () of
        _| abs u < I.unsafeIndex ratios i -> return $! u * bi
         | i == 0                         -> normalTail (u < 0)
         | otherwise                      -> do
             let x  = u * bi
                 xx = x * x
                 d  = exp (-0.5 * (bi * bi - xx))
                 e  = exp (-0.5 * (bj * bj - xx))
             c <- uniform gen
             if e + c * (d - e) < 1
               then return x
               else loop
    blocks = (`I.snoc` 0) . I.cons (v/f) . I.cons r . I.unfoldrN 126 go $! T r f
      where
        go (T b g)   = let !u = T h (exp (-0.5 * h * h))
                           h  = sqrt (-2 * log (v / b + g))
                       in Just (h, u)
        v            = 9.91256303526217e-3
        f            = exp (-0.5 * r * r)
    {-# NOINLINE blocks #-}
    r                = 3.442619855899
    ratios           = I.zipWith (/) (I.tail blocks) blocks
    {-# NOINLINE ratios #-}
    normalTail neg  = tailing
      where tailing  = do
              x <- ((/r) . log) `liftM` uniform gen
              y <- log          `liftM` uniform gen
              if y * (-2) < x * x
                then tailing
                else return $! if neg then x - r else r - x


-- | Generate an exponentially distributed random variate.
exponential :: PrimMonad m
            => Double            -- ^ Scale parameter
            -> Gen (PrimState m) -- ^ Generator
            -> m Double
{-# INLINE exponential #-}<commit_msg>Add workaround proposed by @dzhus for bug #16<commit_after>{-# LANGUAGE BangPatterns #-}
-- |
-- Module    : System.Random.MWC.Distributions
-- Copyright : (c) 2012 <NAME>
-- License   : BSD3
--
-- Maintainer  : <EMAIL>
-- Stability   : experimental
-- Portability : portable
--
-- Pseudo-random number generation for non-uniform distributions.

module System.Random.MWC.Distributions
    (
    -- * Variates: non-uniformly distributed values
      normal
    , standard
    , exponential
    , truncatedExp
    , gamma
    , chiSquare
    , geometric0
    , geometric1

    -- * References
    -- $references
    ) where

import Control.Monad (liftM)
import Control.Monad.Primitive (PrimMonad, PrimState)
import Data.Bits ((.&.))
import Data.Word (Word32)
import System.Random.MWC (Gen, uniform)
import qualified Data.Vector.Unboxed as I

-- Unboxed 2-tuple
data T = T {-# UNPACK #-} !Double {-# UNPACK #-} !Double


-- | Generate a normally distributed random variate with given mean
-- and standard deviation.
normal :: PrimMonad m
       => Double                -- ^ Mean
       -> Double                -- ^ Standard deviation
       -> Gen (PrimState m)
       -> m Double
{-# INLINE normal #-}
-- We express standard in terms of normal and not other way round
-- because of bug in GHC. See bug #16 for more details.
normal m s gen = loop >>= (\x -> return $! m + s * x)
  where
    loop = do
      u  <- (subtract 1 . (*2)) `liftM` uniform gen
      ri <- uniform gen
      let i  = fromIntegral ((ri :: Word32) .&. 127)
          bi = I.unsafeIndex blocks i
          bj = I.unsafeIndex blocks (i+1)
      case () of
        _| abs u < I.unsafeIndex ratios i -> return $! u * bi
         | i == 0                         -> normalTail (u < 0)
         | otherwise                      -> do
             let x  = u * bi
                 xx = x * x
                 d  = exp (-0.5 * (bi * bi - xx))
                 e  = exp (-0.5 * (bj * bj - xx))
             c <- uniform gen
             if e + c * (d - e) < 1
               then return x
               else loop
    blocks = (`I.snoc` 0) . I.cons (v/f) . I.cons r . I.unfoldrN 126 go $! T r f
      where
        go (T b g)   = let !u = T h (exp (-0.5 * h * h))
                           h  = sqrt (-2 * log (v / b + g))
                       in Just (h, u)
        v            = 9.91256303526217e-3
        f            = exp (-0.5 * r * r)
    {-# NOINLINE blocks #-}
    r                = 3.442619855899
    ratios           = I.zipWith (/) (I.tail blocks) blocks
    {-# NOINLINE ratios #-}
    normalTail neg  = tailing
      where tailing  = do
              x <- ((/r) . log) `liftM` uniform gen
              y <- log          `liftM` uniform gen
              if y * (-2) < x * x
                then tailing
                else return $! if neg then x - r else r - x

-- | Generate a normally distributed random variate with zero mean and
-- unit variance.
--
-- The implementation uses Doornik's modified ziggurat algorithm.
-- Compared to the ziggurat algorithm usually used, this is slower,
-- but generates more independent variates that pass stringent tests
-- of randomness.
standard :: PrimMonad m => Gen (PrimState m) -> m Double
{-# INLINE standard #-}
standard = normal 0 1


-- | Generate an exponentially distributed random variate.
exponential :: PrimMonad m
            => Double            -- ^ Scale parameter
            -> Gen (PrimState m) -- ^ Generator
            -> m Double
{-# INLINE exponential #-}
{'fraction_non_alphanumeric': 0.11793002915451894, 'fraction_numerical': 0.019533527696793004, 'mean_word_length': 1.5505576208178438, 'pattern_counts': {'":': 0, '<': 26, '<?xml version=': 0, '>': 34, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '800223', 'n_tokens_mistral': 2209, 'n_tokens_neox': 2029, 'n_words': 725}