Haskell standalone rbinom

rbinom.hs

#!/usr/bin/env cabal
{- cabal:
build-depends: base, mwc-random
ghc-options: -with-rtsopts=-s -Wall
-}

{-
Copyright (C) 1998 Ross Ihaka
Copyright (C) 2000-2014 The R Core Team
Copyright (C) 2007 The R Foundation
Copyright (C) 2019 Jaro Reinders

This program is free software; you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation; either version 2 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program; if not, a copy is available at
https://www.R-project.org/Licenses/

THIS FILE WAS MODIFIED; USE AT YOUR OWN RISK!
-}

{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DeriveFunctor #-}

module Main where

import System.Random.MWC

class UniformRandom m where
  uniformRandom :: m Double

newtype RandIO a = RandIO { unRandIO :: GenIO -> IO a } deriving Functor

instance Applicative RandIO where
  pure x = RandIO (\_ -> return x)
  RandIO p <*> RandIO q = RandIO $ \x -> p x <*> q x

instance Monad RandIO where
  RandIO p >>= f = RandIO $ \x -> p x >>= ($ x) . unRandIO . f

instance UniformRandom RandIO where
  uniformRandom = RandIO uniform

runRandIO :: RandIO a -> IO a
runRandIO = withSystemRandom . unRandIO

main :: IO ()
main = print =<< runRandIO (rbinom 10000000000 0.2)

-- Handle rejection :)
retry :: Monad m => m (Maybe a) -> m a
retry act = maybe (retry act) return =<< act

rbinom :: (UniformRandom m, Monad m) => Int -> Double -> m Int
rbinom n pp = (if pp > 0.5 then fmap (\result -> n - result) else id) $ retry $ do
  u <- fmap (* p4) uniformRandom
  v <- uniformRandom
  return $ case
    if u <= p1 then
      triangular u v
    else if u <= p2 then
      parallelogram u v
    else if u <= p3 then
      leftTail u v
    else
      rightTail u v
    of
      Nothing -> Nothing
      Just (Right x) -> Just x
      Just (Left (v',ix)) -> let k = abs (ix - m) in 
        if k <= 20 || fromIntegral k >= npq / 2 - 1 then
          explicit v' ix
        else
          squeeze k v' ix
  where
    c, fm, npq, p1, p2, p3, p4 :: Double -- , qn :: Double
    xl, xll, xlr, xm, xr :: Double
    m :: Int
    p, q, np, g, r, al0, al, ffm :: Double
    p | pp > 0.5 = 1 - pp
      | otherwise = pp
    q = 1 - p
    np = fromIntegral n * p
    r = p / q
    g = r * fromIntegral (n + 1)
    -- qn = q ^ n

    ffm = np + p
    m = truncate ffm
    fm = fromIntegral m
    npq = np * q
    p1 = fromIntegral (truncate (2.195 * sqrt npq - 4.6 * q) :: Int) + 0.5
    xm = fm + 0.5
    xl = xm - p1
    xr = xm + p1
    c = 0.134 + 20.5 / (15.3 + fm)
    al0 = (ffm - xl) / (ffm - xl * p)
    xll = al0 * (1.0 + 0.5 * al0)
    al = (xr - ffm) / (xr * q)
    xlr = al * (1.0 + 0.5 * al)
    p2 = p1 * (1.0 + c + c)
    p3 = p2 + c / xll
    p4 = p3 + c / xlr

    triangular :: Double -> Double -> Maybe (Either (Double, Int) Int)
    triangular u v = Just (Right (truncate (xm - p1 * v + u)))

    parallelogram :: Double -> Double -> Maybe (Either (Double, Int) Int)
    parallelogram u v
      | v > 1 || v <= 0 = Nothing
      | otherwise = Just (Left (v', truncate x))
      where
        x = xl + (u - p1) / c
        v' = v * c + 1.0 - abs (xm - x) / p1

    leftTail :: Double -> Double -> Maybe (Either (Double, Int) Int)
    leftTail u v 
      | ix < 0 = Nothing
      | otherwise = Just (Left (v * (u - p2) * xll, ix))
      where
        ix = truncate (xl + log v / xll)

    rightTail :: Double -> Double -> Maybe (Either (Double, Int) Int)
    rightTail u v
      | ix > n = Nothing
      | otherwise = Just (Left (v * (u - p3) * xlr, ix))
      where
        ix = truncate (xr - log v / xlr)

    explicit :: Double -> Int -> Maybe Int
    explicit v ix
      | v <= explicit' 1 = Just ix
      | otherwise = Nothing
      where
        explicit' :: Double -> Double
        explicit'
          | m < ix = go1 (m + 1)
          | m /= ix = go2 (ix + 1)
          | otherwise = id

        go1 :: Int -> Double -> Double
        go1 i !f
          | i <= ix = go1 (i + 1) (f * (g / fromIntegral i - r))
          | otherwise = f

        go2 :: Int -> Double -> Double
        go2 i !f
          | i <= m = go2 (i + 1) (f / (g / fromIntegral i - r))
          | otherwise = f

    squeeze :: Int -> Double -> Int -> Maybe Int
    squeeze k v ix
      | alv < ynorm - amaxp = Just ix 
      | alv <= ynorm + amaxp = stirling
      | otherwise = Nothing
      where
        alv, amaxp, ynorm :: Double
        amaxp = (fromIntegral k / npq) * ((fromIntegral k * (fromIntegral k / 3 + 0.625) + 0.1666666666666) / npq + 0.5)
        ynorm = fromIntegral (-k * k) / (2.0 * npq)
        alv = log v
        
        stirling :: Maybe Int
        stirling
          | alv <= xm * log (f1 / x1) + (fromIntegral (n - m) + 0.5) * log (z / w) + fromIntegral (ix - m) * log (w * p / (x1 * q)) + (13860.0 - (462.0 - (132.0 - (99.0 - 140.0 / f2) / f2) / f2) / f2) / f1 / 166320.0 + (13860.0 - (462.0 - (132.0 - (99.0 - 140.0 / z2) / z2) / z2) / z2) / z / 166320.0 + (13860.0 - (462.0 - (132.0 - (99.0 - 140.0 / x2) / x2) / x2) / x2) / x1 / 166320.0 + (13860.0 - (462.0 - (132.0 - (99.0 - 140.0 / w2) / w2) / w2) / w2) / w / 166320.0 = Just ix
          | otherwise = Nothing
          where
            f1, f2, w, w2, x1, x2, z, z2 :: Double
            x1 = fromIntegral ix + 1
            f1 = fm + 1.0
            z = fromIntegral (n + 1) - fm
            w = fromIntegral (n - ix) + 1.0
            z2 = z * z
            x2 = x1 * x1
            f2 = f1 * f1
            w2 = w * w