Haskell infinite list of Bernoulli distributed booleans - haskell

Question

I need an list of biased, random booleans. Each boolean needs to have the same probability of being True (Bernoulli distributed). These booleans are passed to a function, which generates zero or more output booleans per input boolean. I need an infinite list, because I don't know in advance how many booleans are required to provide enough output. See the below (simplified) code:
import System.Random.MWC
import System.Random.MWC.Distributions
foo :: [Bool] -> [Bool] -- foo outputs zero or more Bools per input Bool
main = do
gen <- create
bits <- sequence . repeat $ bernoulli 0.25 gen
print . take 32 . foo $ bits
Unfortunately, this code just hangs at the second line of main. I guess that there is something non-lazy happening somewhere with Control.Monad.ST?
(I would be able to do something like this with System.Random.randoms, but the resulting values don't have the required distributions.)
Can I fix this while keep using the System.Random.MWC library? Or does this require me to switch to alternative implementations?

Answer 1

The mwc-random package provides two PrimMonad instances, one for IO and another for ST s. As long as an ST computation is parameterized over all state tags s, we can run the computation and extract the value with runST :: (forall s. ST s a) -> a. By itself this wouldn't be very useful since we'd lose the state: the seed of the random generator, but mwc-random also provides explicit ways to handle the seeds:
save :: PrimMonad m => Gen (PrimState m) -> m Seed
restore :: PrimMonad m => Seed -> m (Gen (PrimState m))
We can use these to make a computation that generates a stream of values from a computation that generates a single value, as long as the generator is in forall s. ST s.
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
import System.Random.MWC
import Control.Monad.ST
import System.Random.MWC.Distributions
randomStream :: forall s a. (forall s. GenST s -> ST s a) -> GenST s -> ST s [a]
randomStream item = go
where
go :: forall s. GenST s -> ST s [a]
go gen = do
x <- item gen
seed <- save gen
return (x:runST (restore seed >>= go))
With this we can write your example as
main = do
bits <- withSystemRandom (randomStream (bernoulli 0.25))
print . take 32 $ bits
We can actually build generators more sophisticated than using the same generator for each item in the stream. We could thread a state along the stream so that each value can depend on the results of the previous ones.
unfoldStream :: forall s a b. (forall s. b -> GenST s -> ST s (a, b)) -> b -> GenST s -> ST s [a]
unfoldStream item = go
where
go :: forall s. b -> GenST s -> ST s [a]
go b gen = do
(x,b') <- item b gen
seed <- save gen
return (x:runST (restore seed >>= go b'))
The following example stream has results that increase in likelihood every time the result is False.
import Control.Monad.Primitive
interesting :: (PrimMonad m) => Double -> Gen (PrimState m) -> m (Bool, Double)
interesting p gen = do
result <- bernoulli p gen
let p' = if result then p else p + (1-p)*0.25
return (result, p')
main = do
bits <- withSystemRandom (unfoldStream interesting 0)
print . take 32 $ bits

Answer 2

The culprit is sequence . repeat - this will hang for (almost?) every monad, since you must perform a potentially infinite number of effects.
The simplest solutions would be to use a different library - which may not be possible if you are relying on the quality of the numbers produced from mwc-random. The next simplest solution is to rewrite foo to have type [IO Bool] -> IO [Bool] and pass it repeat (bernoulli 0.25 gen) - this would allow foo to make the choice of when to stop executing the effects produced by the infinite list. But having your logic inside of IO is not very nice.
The standard trick when you need an infinite list of random numbers is to use a pure function f :: StdGen -> (Result, StdGen). Then unfoldr (Just . f) :: StdGen -> [Result], and the output is an infinite list. At first glance, it may appear that mwc-random only has monadic functions, and that there is no pure interface. However, that is not the case, because ST s is an instance of PrimMonad. You also have the functions converting a Gen to a Seed. Using these, you can get a pure RNG function for any monadic one:
{-# LANGUAGE RankNTypes #-}
import System.Random.MWC
import System.Random.MWC.Distributions
import Control.Monad.ST
import Data.List
pureRand :: (forall s . GenST s -> ST s t) -> Seed -> (t, Seed)
pureRand f s = runST $ do
s' <- restore s
r <- f s'
s'' <- save s'
return (r, s'')
pureBernoulli :: Double -> Seed -> (Bool, Seed)
pureBernoulli a = pureRand (bernoulli a)
foo :: [Bool] -> [Bool]
foo = id
main = do
gen <- create >>= save
let bits = unfoldr (Just . pureBernoulli 0.25) gen
print . take 32 . foo $ bits
It is unfortunate that mwc-random doesn't expose this sort of interface by default but it is pretty easy to get to.
The other option is slightly more scary - use unsafe functions.
import System.IO.Unsafe
repeatM rand = go where
go = do
x <- rand
xs <- unsafeInterleaveIO go
return (x : xs)
main2 = do
gen <- create
bits <- repeatM (bernoulli 0.25 gen)
print . take 32 . foo $ bits
Naturally this comes with the usual caveats surrounding unsafe - use it only if you are exceedingly inconvenienced by the pure functions. unsafeInterleaveIO may reorder or never execute effects - if foo, for example, ignores one element, it will never be computed and the corresponding effect of updating the state stored in gen may not happen. For example, the following will print nothing:
snd <$> ((,) <$> unsafeInterleaveIO (putStrLn "Hello") <*> return ())