Box-Muller for MonadRandom

BoxMuller.hs

module BoxMuller
  ( getNormal
  , getNormal'
  , getNormals
  , getNormals'
  ) where

import Control.Monad ((<$!>))
import Control.Monad.Random

-- | Sample a random number from N(0,1).
--
-- The Box-Muller algorithm generates numbers in pairs. If you need more than
-- one, consider using
--
-- @a:b:_ \<- getNormals@
--
-- or
--
-- @take n \<$\> getNormals@
getNormal :: (MonadRandom m, Floating n, Random n) => m n
getNormal = (\ ~(n:_) -> n) <$!> getNormals

-- | Sample a random number from N(mean,sigma).
--
-- The Box-Muller algorithm generates numbers in pairs. If you need more than
-- one, consider using
--
-- @a:b:_ \<- getNormals' mean sigma@
--
-- or
--
-- @take n \<$\> getNormals' mean sigma@
getNormal' :: (MonadRandom m, Floating n, Random n)
           => n  -- ^ mean
           -> n  -- ^ sigma
           -> m n
getNormal' mean sigma = (\ ~(n:_) -> n) <$!> getNormals' mean sigma

-- | Sample an infinite list of random numbers from N(0,1).
getNormals :: (MonadRandom m, Floating n, Random n) => m [n]
getNormals = (\(a,b) xs -> a:b:xs) <$> boxMuller <*> getNormals

-- | Sample an infinite list of random numbers from N(mean,sigma).
getNormals' :: (MonadRandom m, Floating n, Random n)
            => n  -- ^ mean
            -> n  -- ^ sigma
            -> m [n]
getNormals' mean sigma = map (\n -> mean + sigma * n) <$> getNormals

-- http://en.wikipedia.org/wiki/Box%E2%80%93Muller_transform
boxMuller :: (MonadRandom m, Floating n, Random n) => m (n,n)
boxMuller = go <$> getRandomR (0,1) <*> getRandomR (0,1)
  where
    go u0 u1 = z0 `seq` z1 `seq` (z0,z1)
      where z0 = r * cos phi
            z1 = r * sin phi
            r = sqrt (-2 * log u0)
            phi = 2 * pi * u1